Simplify input transform object (#1340)

docs/source/reference/image_classification.rst

@@ -108,6 +108,7 @@ Here's an example:
     from typing import Callable, Tuple, Union
     import flash
     from flash.image import ImageClassificationData, ImageClassifier
+    from flash.core.data.transforms import ApplyToKeys
     from flash.core.data.io.input_transform import InputTransform
     from dataclasses import dataclass
 
@@ -119,24 +120,37 @@ Here's an example:
         mean: Union[float, Tuple[float, float, float]] = (0.485, 0.456, 0.406)
         std: Union[float, Tuple[float, float, float]] = (0.229, 0.224, 0.225)
 
-        def input_per_sample_transform(self):
-            return T.Compose([T.ToTensor(), T.Resize(self.image_size), T.Normalize(self.mean, self.std)])
-
-        def train_input_per_sample_transform(self):
+        def per_sample_transform(self):
             return T.Compose(
                 [
-                    T.ToTensor(),
-                    T.Resize(self.image_size),
-                    T.Normalize(self.mean, self.std),
-                    T.RandomHorizontalFlip(),
-                    T.ColorJitter(),
-                    T.RandomAutocontrast(),
-                    T.RandomPerspective(),
+                    ApplyToKeys(
+                        "input",
+                        T.Compose([T.ToTensor(), T.Resize(self.image_size), T.Normalize(self.mean, self.std)]),
+                    ),
+                    ApplyToKeys("target", torch.as_tensor),
                 ]
             )
 
-        def target_per_sample_transform(self) -> Callable:
-            return torch.as_tensor
+        def train_per_sample_transform(self):
+            return T.Compose(
+                [
+                    ApplyToKeys(
+                        "input",
+                        T.Compose(
+                            [
+                                T.ToTensor(),
+                                T.Resize(self.image_size),
+                                T.Normalize(self.mean, self.std),
+                                T.RandomHorizontalFlip(),
+                                T.ColorJitter(),
+                                T.RandomAutocontrast(),
+                                T.RandomPerspective(),
+                            ]
+                        ),
+                    ),
+                    ApplyToKeys("target", torch.as_tensor),
+                ]
+            )
 
 
     datamodule = ImageClassificationData.from_folders(

docs/source/template/data.rst

@@ -80,28 +80,20 @@ Internally we inject the :class:`~flash.core.data.io.input_transform.InputTransf
 
 Defining the standard transforms (typically at least a ``per_sample_transform`` should be defined) for your :class:`~flash.core.data.io.input_transform.InputTransform` involves simply overriding the required hook to return a callable transform.
 
-For our ``TemplateInputTransform``, we'll just configure an ``input_per_sample_transform`` and a ``target_per_sample_transform``.
+For our ``TemplateInputTransform``, we'll just configure a ``per_sample_transform``.
 Let's first define a to_tensor transform as a ``staticmethod``:
 
 .. literalinclude:: ../../../flash/template/classification/data.py
     :language: python
     :dedent: 4
-    :pyobject: TemplateInputTransform.input_to_tensor
+    :pyobject: TemplateInputTransform.to_tensor
 
-Now in our ``input_per_sample_transform`` hook, we return the transform:
+Now in our ``per_sample_transform`` hook, we return the transform:
 
 .. literalinclude:: ../../../flash/template/classification/data.py
     :language: python
     :dedent: 4
-    :pyobject: TemplateInputTransform.input_per_sample_transform
-
-To convert the targets to a tensor we can simply use ``torch.as_tensor``.
-Here's our ``target_per_sample_transform``:
-
-.. literalinclude:: ../../../flash/template/classification/data.py
-    :language: python
-    :dedent: 4
-    :pyobject: TemplateInputTransform.target_per_sample_transform
+    :pyobject: TemplateInputTransform.per_sample_transform
 
 .. _contributing_data_module:
 

flash/audio/classification/input_transform.py

@@ -16,7 +16,9 @@
 
 import torch
 
+from flash.core.data.io.input import DataKeys
 from flash.core.data.io.input_transform import InputTransform
+from flash.core.data.transforms import ApplyToKeys
 from flash.core.utilities.imports import _TORCHAUDIO_AVAILABLE, _TORCHVISION_AVAILABLE
 
 if _TORCHVISION_AVAILABLE:
@@ -33,7 +35,7 @@ class AudioClassificationInputTransform(InputTransform):
     time_mask_param: Optional[int] = None
     freq_mask_param: Optional[int] = None
 
-    def train_input_per_sample_transform(self) -> Callable:
+    def train_per_sample_transform(self) -> Callable:
         transforms = []
         if self.time_mask_param is not None:
             transforms.append(TAudio.TimeMasking(time_mask_param=self.time_mask_param))
@@ -42,10 +44,17 @@ def train_input_per_sample_transform(self) -> Callable:
             transforms.append(TAudio.FrequencyMasking(freq_mask_param=self.freq_mask_param))
 
         transforms += [T.ToTensor(), T.Resize(self.spectrogram_size)]
-        return T.Compose(transforms)
-
-    def input_per_sample_transform(self) -> Callable:
-        return T.Compose([T.ToTensor(), T.Resize(self.spectrogram_size)])
-
-    def target_per_sample_transform(self) -> Callable:
-        return torch.as_tensor
+        return T.Compose(
+            [
+                ApplyToKeys(DataKeys.INPUT, T.Compose(transforms)),
+                ApplyToKeys(DataKeys.TARGET, torch.as_tensor),
+            ]
+        )
+
+    def per_sample_transform(self) -> Callable:
+        return T.Compose(
+            [
+                ApplyToKeys(DataKeys.INPUT, T.Compose([T.ToTensor(), T.Resize(self.spectrogram_size)])),
+                ApplyToKeys(DataKeys.TARGET, torch.as_tensor),
+            ]
+        )