Repository: Lightning-Universe/lightning-flash

Issue:
Support for a list of transforms rather than dict
## 🚀 Feature

Enable support for just a list of transforms / a single Sequential or Compose transform to be passed to the preprocess. This could just be bound to the `pre_tensor_transform` hook by default.



Discussion:
Hi @ethanwharris thank you for taking your time to open this issue. I think the current way is not bad but it should be documented better :) 
Hi @ethanwharris, it would be awesome to have this feature. Also, we can have a method to get augmentations like below - 

```python
def get_augs(rotate_vertical:bool, rotate_horizontal:bool, shear:float):
     """do something"""
    return augmentations
```
Then this can be passed directly into `from_folder` or `from_df` method.

keras has [ImageDataGenerator](https://keras.io/api/preprocessing/image/) API which makes it easier to use for new to advance users.
@aniketmaurya Want to try to add it? I think we can just detect if transforms are not a dict but instead a list or callable then wrap in a dictionary. I like your suggestion of a `get_augs`, an idea we had is to make a registry for transform recipes, then we could just add a bunch of helpful defaults that can be referenced with a string passed as `*_transform`. That has the advantage of working nicely with Flash Zero. What do you think? cc @ananyahjha93 who has worked on something related to this
Sure @ethanwharris I will create a PR for the first task first - support transformation without dict.

Resolving patch:
diff --git a/flash/core/data/data_module.py b/flash/core/data/data_module.py
--- a/flash/core/data/data_module.py
+++ b/flash/core/data/data_module.py
@@ -460,9 +460,9 @@ def from_data_source(
         val_data: Any = None,
         test_data: Any = None,
         predict_data: Any = None,
-        train_transform: Optional[Dict[str, Callable]] = None,
-        val_transform: Optional[Dict[str, Callable]] = None,
-        test_transform: Optional[Dict[str, Callable]] = None,
+        train_transform: Optional[Union[Callable, List, Dict[str, Callable]]] = None,
+        val_transform: Optional[Union[Callable, List, Dict[str, Callable]]] = None,
+        test_transform: Optional[Union[Callable, List, Dict[str, Callable]]] = None,
         predict_transform: Optional[Dict[str, Callable]] = None,
         data_fetcher: Optional[BaseDataFetcher] = None,
         preprocess: Optional[Preprocess] = None,
@@ -522,6 +522,7 @@ def from_data_source(
                 },
             )
         """
+
         preprocess = preprocess or cls.preprocess_cls(
             train_transform,
             val_transform,
diff --git a/flash/core/data/process.py b/flash/core/data/process.py
--- a/flash/core/data/process.py
+++ b/flash/core/data/process.py
@@ -14,7 +14,7 @@
 import inspect
 import os
 from abc import ABC, abstractclassmethod, abstractmethod
-from typing import Any, Callable, Dict, List, Mapping, Optional, Sequence
+from typing import Any, Callable, Dict, List, Mapping, Optional, Sequence, Union
 
 import torch
 from pytorch_lightning.trainer.states import RunningStage
@@ -28,6 +28,7 @@
 from flash.core.data.data_source import DatasetDataSource, DataSource, DefaultDataKeys, DefaultDataSources
 from flash.core.data.properties import Properties
 from flash.core.data.states import CollateFn
+from flash.core.data.transforms import ApplyToKeys
 from flash.core.data.utils import _PREPROCESS_FUNCS, _STAGES_PREFIX, convert_to_modules, CurrentRunningStageFuncContext
 
 
@@ -177,10 +178,10 @@ def pre_tensor_transform(self, sample: PIL.Image) -> PIL.Image:
 
     def __init__(
         self,
-        train_transform: Optional[Dict[str, Callable]] = None,
-        val_transform: Optional[Dict[str, Callable]] = None,
-        test_transform: Optional[Dict[str, Callable]] = None,
-        predict_transform: Optional[Dict[str, Callable]] = None,
+        train_transform: Optional[Union[Callable, List, Dict[str, Callable]]] = None,
+        val_transform: Optional[Union[Callable, List, Dict[str, Callable]]] = None,
+        test_transform: Optional[Union[Callable, List, Dict[str, Callable]]] = None,
+        predict_transform: Optional[Union[Callable, List, Dict[str, Callable]]] = None,
         data_sources: Optional[Dict[str, "DataSource"]] = None,
         deserializer: Optional["Deserializer"] = None,
         default_data_source: Optional[str] = None,
@@ -252,6 +253,11 @@ def _check_transforms(
         if transform is None:
             return transform
 
+        if isinstance(transform, list):
+            transform = {"pre_tensor_transform": ApplyToKeys(DefaultDataKeys.INPUT, torch.nn.Sequential(*transform))}
+        elif callable(transform):
+            transform = {"pre_tensor_transform": ApplyToKeys(DefaultDataKeys.INPUT, transform)}
+
         if not isinstance(transform, Dict):
             raise MisconfigurationException(
                 "Transform should be a dict. " f"Here are the available keys for your transforms: {_PREPROCESS_FUNCS}."
@@ -439,10 +445,10 @@ def data_source_of_name(self, data_source_name: str) -> DataSource:
 class DefaultPreprocess(Preprocess):
     def __init__(
         self,
-        train_transform: Optional[Dict[str, Callable]] = None,
-        val_transform: Optional[Dict[str, Callable]] = None,
-        test_transform: Optional[Dict[str, Callable]] = None,
-        predict_transform: Optional[Dict[str, Callable]] = None,
+        train_transform: Optional[Union[Callable, List, Dict[str, Callable]]] = None,
+        val_transform: Optional[Union[Callable, List, Dict[str, Callable]]] = None,
+        test_transform: Optional[Union[Callable, List, Dict[str, Callable]]] = None,
+        predict_transform: Optional[Union[Callable, List, Dict[str, Callable]]] = None,
         data_sources: Optional[Dict[str, "DataSource"]] = None,
         default_data_source: Optional[str] = None,
     ):
diff --git a/flash/image/classification/data.py b/flash/image/classification/data.py
--- a/flash/image/classification/data.py
+++ b/flash/image/classification/data.py
@@ -120,9 +120,9 @@ def from_data_frame(
         predict_data_frame: Optional[pd.DataFrame] = None,
         predict_images_root: Optional[str] = None,
         predict_resolver: Optional[Callable[[str, str], str]] = None,
-        train_transform: Optional[Dict[str, Callable]] = None,
-        val_transform: Optional[Dict[str, Callable]] = None,
-        test_transform: Optional[Dict[str, Callable]] = None,
+        train_transform: Optional[Union[Callable, List, Dict[str, Callable]]] = None,
+        val_transform: Optional[Union[Callable, List, Dict[str, Callable]]] = None,
+        test_transform: Optional[Union[Callable, List, Dict[str, Callable]]] = None,
         predict_transform: Optional[Dict[str, Callable]] = None,
         data_fetcher: Optional[BaseDataFetcher] = None,
         preprocess: Optional[Preprocess] = None,
@@ -217,9 +217,9 @@ def from_csv(
         predict_file: Optional[str] = None,
         predict_images_root: Optional[str] = None,
         predict_resolver: Optional[Callable[[str, str], str]] = None,
-        train_transform: Optional[Dict[str, Callable]] = None,
-        val_transform: Optional[Dict[str, Callable]] = None,
-        test_transform: Optional[Dict[str, Callable]] = None,
+        train_transform: Optional[Union[Callable, List, Dict[str, Callable]]] = None,
+        val_transform: Optional[Union[Callable, List, Dict[str, Callable]]] = None,
+        test_transform: Optional[Union[Callable, List, Dict[str, Callable]]] = None,
         predict_transform: Optional[Dict[str, Callable]] = None,
         data_fetcher: Optional[BaseDataFetcher] = None,
         preprocess: Optional[Preprocess] = None,


Test patch:
diff --git a/tests/core/data/test_process.py b/tests/core/data/test_process.py
--- a/tests/core/data/test_process.py
+++ b/tests/core/data/test_process.py
@@ -88,7 +88,6 @@ class Serializer2State(ProcessState):
 
 
 def test_saving_with_serializers(tmpdir):
-
     checkpoint_file = os.path.join(tmpdir, "tmp.ckpt")
 
     class CustomModel(Task):
@@ -122,7 +121,6 @@ def __init__(self):
 
 
 def test_data_source_of_name():
-
     preprocess = CustomPreprocess()
 
     assert preprocess.data_source_of_name("test")() == "test"
@@ -135,7 +133,6 @@ def test_data_source_of_name():
 
 
 def test_available_data_sources():
-
     preprocess = CustomPreprocess()
 
     assert DefaultDataSources.TENSORS in preprocess.available_data_sources()
@@ -147,3 +144,9 @@ def test_available_data_sources():
     assert DefaultDataSources.TENSORS in data_module.available_data_sources()
     assert "test" in data_module.available_data_sources()
     assert len(data_module.available_data_sources()) == 3
+
+
+def test_check_transforms():
+    transform = torch.nn.Identity()
+    DefaultPreprocess(train_transform=transform)
+    DefaultPreprocess(train_transform=[transform])