Document custom image transformations

docs/source/reference/image_classification.rst

@@ -57,6 +57,57 @@ Here's the full example:
 
 ------
 
+**********************
+Custom Transformations
+**********************
+
+Flash automatically applies some default image transformations and augmentations, but you may wish to customize these for your own use case.
+The base :class:`~flash.core.data.process.Preprocess` defines 7 hooks for different stages in the data loading pipeline.
+To apply image augmentations you can directly import the ``default_transforms`` from ``flash.image.classification.transforms`` and then merge your custom image transformations with them using the :func:`~flash.core.data.transforms.merge_transforms` helper function.
+Here's an example where we load the default transforms and merge with custom `torchvision` transformations.
+We use the `post_tensor_transform` hook to apply the transformations after the image has been converted to a `torch.Tensor`.
+
+
+.. testsetup:: transformations
+
+    from flash.core.data.utils import download_data
+
+    download_data("https://pl-flash-data.s3.amazonaws.com/hymenoptera_data.zip", "./data")
+
+.. testcode:: transformations
+
+    from torchvision import transforms as T
+
+    import flash
+    from flash.core.data.data_source import DefaultDataKeys
+    from flash.core.data.transforms import ApplyToKeys, merge_transforms
+    from flash.image import ImageClassificationData, ImageClassifier
+    from flash.image.classification.transforms import default_transforms
+
+    post_tensor_transform = ApplyToKeys(
+        DefaultDataKeys.INPUT,
+        T.Compose([T.RandomHorizontalFlip(), T.ColorJitter(), T.RandomAutocontrast(), T.RandomPerspective()]),
+    )
+
+    new_transforms = merge_transforms(default_transforms((64, 64)), {"post_tensor_transform": post_tensor_transform})
+
+    datamodule = ImageClassificationData.from_folders(
+        train_folder="data/hymenoptera_data/train/", val_folder="data/hymenoptera_data/val/", train_transform=new_transforms
+    )
+
+    model = ImageClassifier(backbone="resnet18", num_classes=datamodule.num_classes)
+
+    trainer = flash.Trainer(max_epochs=1)
+    trainer.finetune(model, datamodule=datamodule, strategy="freeze")
+
+
+.. testoutput:: transformations
+    :hide:
+
+    ...
+
+------
+
 **********
 Flash Zero
 **********