@@ -4,20 +4,36 @@ from omegaconf import ListConfig
 
 from super_gradients.common.factories.base_factory import BaseFactory
 from super_gradients.common.factories.list_factory import ListFactory
-from super_gradients.training.transforms import TRANSFORMS
+from super_gradients.training.transforms import TRANSFORMS, ALBUMENTATIONS_TRANSFORMS, ALBUMENTATIONS_COMP_TRANSFORMS, imported_albumentations_failure
+from super_gradients.training.transforms.pipeline_adaptors import AlbumentationsAdaptor
 
 
 class TransformsFactory(BaseFactory):
-
     def __init__(self):
         super().__init__(TRANSFORMS)
 
     def get(self, conf: Union[str, dict]):
 
-        # SPECIAL HANDLING FOR COMPOSE
-        if isinstance(conf, Mapping) and 'Compose' in conf:
-            conf['Compose']['transforms'] = ListFactory(TransformsFactory()).get(conf['Compose']['transforms'])
+        # SPECIAL HANDLING FOR COMPOSE AND ALBUMENTATIONS
+        if isinstance(conf, Mapping) and "Albumentations" in conf:
+            return AlbumentationsAdaptor(AlbumentationsTransformsFactory().get(conf["Albumentations"]))
+        if isinstance(conf, Mapping) and "Compose" in conf:
+            conf["Compose"]["transforms"] = ListFactory(TransformsFactory()).get(conf["Compose"]["transforms"])
         elif isinstance(conf, (list, ListConfig)):
             conf = ListFactory(TransformsFactory()).get(conf)
 
         return super().get(conf)
+
+
+class AlbumentationsTransformsFactory(BaseFactory):
+    def __init__(self):
+        if imported_albumentations_failure:
+            raise imported_albumentations_failure
+        super().__init__(ALBUMENTATIONS_TRANSFORMS)
+
+    def get(self, conf: Union[str, dict]):
+        if isinstance(conf, Mapping):
+            _type = list(conf.keys())[0]  # THE TYPE NAME
+            if _type in ALBUMENTATIONS_COMP_TRANSFORMS:
+                conf[_type]["transforms"] = ListFactory(AlbumentationsTransformsFactory()).get(conf[_type]["transforms"])
+        return super(AlbumentationsTransformsFactory, self).get(conf)

@@ -5,7 +5,9 @@
 #   1. Move to the project root (where you will find the ReadMe and src folder)
 #   2. Run the command:
 #       python src/super_gradients/examples/train_from_recipe_example/train_from_recipe.py --config-name=cifar10_resnet +experiment_name=cifar10
-
+#
+#   To use equivalent Albumentations transforms pipeline set dataset_params to cifar10_albumentations_dataset_params:
+#     python src/super_gradients/examples/train_from_recipe_example/train_from_recipe.py --config-name=cifar10_resnet dataset_params=cifar10_albumentations_dataset_params
 defaults:
   - training_hyperparams: cifar10_resnet_train_params
   - dataset_params: cifar10_dataset_params

# Equivalent to cifar10_dataset_params.yaml, but uses albumentations transforms.
# The purpose of the below configuration is to demonstrate the use of Albumentation transforms in train_from_recipe.

batch_size: 256 # batch size for trainset
val_batch_size: 512 # batch size for valset in DatasetInterface

# TODO: REMOVE ABOVE, HERE FOR COMPATIBILITY UNTIL WE REMOVE DATASET_INTERFACE

train_dataset_params:
  root: ./data/cifar10
  train: True
  transforms:
    Albumentations:
      Compose:
        transforms:
          - RandomCrop:
              height: 32
              width: 32
          - HorizontalFlip:
              p: 0.5
          - Normalize:
              mean:
                - 0.4914
                - 0.4822
                - 0.4465
              std:
                - 0.2023
                - 0.1994
                - 0.2010
          - ToTensorV2
  target_transform: null
  download: True

train_dataloader_params:
  batch_size: 256
  num_workers: 8
  drop_last: False
  pin_memory: True

val_dataset_params:
  root: ./data/cifar10
  train: False
  transforms:
    Albumentations:
      Compose:
        transforms:
          - Normalize:
              mean:
                - 0.4914
                - 0.4822
                - 0.4465
              std:
                - 0.2023
                - 0.1994
                - 0.2010
          - ToTensorV2
  target_transform: null
  download: True

val_dataloader_params:
  batch_size: 512
  num_workers: 8
  drop_last: False
  pin_memory: True

@@ -1,8 +1,7 @@
-from typing import Optional, Callable
+from typing import Optional, Callable, Union
 
 from torchvision.transforms import Compose
 
-from super_gradients.common.factories.list_factory import ListFactory
 from super_gradients.common.factories.transforms_factory import TransformsFactory
 from super_gradients.common.decorators.factory_decorator import resolve_param
 from torchvision.datasets import CIFAR10, CIFAR100
@@ -18,31 +17,37 @@ class Cifar10(CIFAR10):
     :param target_transform:        Transform to apply to target output
     :param download:                Download (True) the dataset from source
     """
-    @resolve_param("transforms", ListFactory(TransformsFactory()))
+
+    @resolve_param("transforms", TransformsFactory())
     def __init__(
         self,
         root: str,
         train: bool = True,
-        transforms: Optional[Callable] = None,
+        transforms: Union[list, dict] = None,
         target_transform: Optional[Callable] = None,
         download: bool = False,
     ) -> None:
+        # TO KEEP BACKWARD COMPATABILITY, WILL BE REMOVED IN THE FUTURE ONCE WE ALLIGN TORCHVISION/NATIVE TRANSFORMS
+        # TREATMENT IN FACTORIES (I.E STATING COMPOSE IN CONFIGS)
+        if isinstance(transforms, list):
+            transforms = Compose(transforms)
+
         super(Cifar10, self).__init__(
             root=root,
             train=train,
-            transform=Compose(transforms),
+            transform=transforms,
             target_transform=target_transform,
             download=download,
         )
 
 
 class Cifar100(CIFAR100):
-    @resolve_param("transforms", ListFactory(TransformsFactory()))
+    @resolve_param("transforms", TransformsFactory())
     def __init__(
         self,
         root: str,
         train: bool = True,
-        transforms: Optional[Callable] = None,
+        transforms: Union[list, dict] = None,
         target_transform: Optional[Callable] = None,
         download: bool = False,
     ) -> None:
@@ -55,10 +60,15 @@ class Cifar100(CIFAR100):
         :param target_transform:        Transform to apply to target output
         :param download:                Download (True) the dataset from source
         """
+        # TO KEEP BACKWARD COMPATABILITY, WILL BE REMOVED IN THE FUTURE ONCE WE ALLIGN TORCHVISION/NATIVE TRANSFORMS
+        # TREATMENT IN FACTORIES (I.E STATING COMPOSE IN CONFIGS)
+        if isinstance(transforms, list):
+            transforms = Compose(transforms)
+
         super(Cifar100, self).__init__(
             root=root,
             train=train,
-            transform=Compose(transforms),
+            transform=transforms,
             target_transform=target_transform,
             download=download,
         )

@@ -1,14 +1,19 @@
+from typing import Union
+
 import torchvision.datasets as torch_datasets
 from torchvision.transforms import Compose
 
 from super_gradients.common.decorators.factory_decorator import resolve_param
 from super_gradients.common.factories.transforms_factory import TransformsFactory
-from super_gradients.common.factories.list_factory import ListFactory
 
 
 class ImageNetDataset(torch_datasets.ImageFolder):
     """ImageNetDataset dataset"""
 
-    @resolve_param('transforms', factory=ListFactory(TransformsFactory()))
-    def __init__(self, root: str, transforms: list = [], *args, **kwargs):
-        super(ImageNetDataset, self).__init__(root, transform=Compose(transforms), *args, **kwargs)
+    @resolve_param("transforms", factory=TransformsFactory())
+    def __init__(self, root: str, transforms: Union[list, dict] = [], *args, **kwargs):
+        # TO KEEP BACKWARD COMPATABILITY, WILL BE REMOVED IN THE FUTURE ONCE WE ALLIGN TORCHVISION/NATIVE TRANSFORMS
+        # TREATMENT IN FACTORIES (I.E STATING COMPOSE IN CONFIGS)
+        if isinstance(transforms, list):
+            transforms = Compose(transforms)
+        super(ImageNetDataset, self).__init__(root, transform=transforms, *args, **kwargs)

@@ -1,10 +1,31 @@
 # PACKAGE IMPORTS FOR EXTERNAL USAGE
 import cv2
-from super_gradients.training.transforms.transforms import DetectionMosaic, DetectionRandomAffine, DetectionHSV,\
-    DetectionPaddedRescale, DetectionTargetsFormatTransform
-from super_gradients.training.transforms.all_transforms import TRANSFORMS, Transforms
+from super_gradients.training.transforms.transforms import (
+    DetectionMosaic,
+    DetectionRandomAffine,
+    DetectionHSV,
+    DetectionPaddedRescale,
+    DetectionTargetsFormatTransform,
+)
+from super_gradients.training.transforms.all_transforms import (
+    TRANSFORMS,
+    ALBUMENTATIONS_TRANSFORMS,
+    Transforms,
+    ALBUMENTATIONS_COMP_TRANSFORMS,
+    imported_albumentations_failure,
+)
 
-__all__ = ['TRANSFORMS', 'Transforms', 'DetectionMosaic', 'DetectionRandomAffine', 'DetectionHSV', 'DetectionPaddedRescale',
-           'DetectionTargetsFormatTransform']
+__all__ = [
+    "TRANSFORMS",
+    "ALBUMENTATIONS_TRANSFORMS",
+    "ALBUMENTATIONS_COMP_TRANSFORMS",
+    "Transforms",
+    "DetectionMosaic",
+    "DetectionRandomAffine",
+    "DetectionHSV",
+    "DetectionPaddedRescale",
+    "DetectionTargetsFormatTransform",
+    "imported_albumentations_failure",
+]
 
 cv2.setNumThreads(0)

@@ -1,6 +1,12 @@
 from super_gradients.common.object_names import Transforms
 from super_gradients.training.datasets.data_augmentation import Lighting, RandomErase
 from super_gradients.training.datasets.datasets_utils import RandomResizedCropAndInterpolation, rand_augment_transform
+import importlib
+import inspect
+
+from super_gradients.common.abstractions.abstract_logger import get_logger
+
+
 from super_gradients.training.transforms.transforms import (
     SegRandomFlip,
     SegRescale,
@@ -118,3 +124,31 @@ TRANSFORMS = {
     Transforms.RandomAutocontrast: RandomAutocontrast,
     Transforms.RandomEqualize: RandomEqualize,
 }
+logger = get_logger(__name__)
+
+try:
+    from albumentations import BasicTransform, BaseCompose
+
+    imported_albumentations_failure = None
+except (ImportError, NameError, ModuleNotFoundError) as import_err:
+    logger.debug("Failed to import pytorch_quantization")
+    imported_albumentations_failure = import_err
+
+if imported_albumentations_failure is None:
+    ALBUMENTATIONS_TRANSFORMS = {
+        name: cls for name, cls in inspect.getmembers(importlib.import_module("albumentations"), inspect.isclass) if issubclass(cls, BasicTransform)
+    }
+    ALBUMENTATIONS_TRANSFORMS.update(
+        {name: cls for name, cls in inspect.getmembers(importlib.import_module("albumentations.pytorch"), inspect.isclass) if issubclass(cls, BasicTransform)}
+    )
+
+    ALBUMENTATIONS_COMP_TRANSFORMS = {
+        name: cls
+        for name, cls in inspect.getmembers(importlib.import_module("albumentations.core.composition"), inspect.isclass)
+        if issubclass(cls, BaseCompose)
+    }
+    ALBUMENTATIONS_TRANSFORMS.update(ALBUMENTATIONS_COMP_TRANSFORMS)
+
+else:
+    ALBUMENTATIONS_TRANSFORMS = None
+    ALBUMENTATIONS_COMP_TRANSFORMS = None

from typing import Callable
from abc import abstractmethod, ABC
import numpy as np


class TransformsPipelineAdaptorBase(ABC):
    def __init__(self, composed_transforms: Callable):
        self.composed_transforms = composed_transforms

    @abstractmethod
    def __call__(self, sample, *args, **kwargs):
        raise NotImplementedError

    @abstractmethod
    def prep_for_transforms(self, sample):
        raise NotImplementedError

    @abstractmethod
    def post_transforms_processing(self, sample):
        raise NotImplementedError


class AlbumentationsAdaptor(TransformsPipelineAdaptorBase):
    def __init__(self, composed_transforms: Callable):
        super(AlbumentationsAdaptor, self).__init__(composed_transforms)

    def __call__(self, sample, *args, **kwargs):
        sample = self.prep_for_transforms(sample)
        sample = self.composed_transforms(**sample)["image"]
        sample = self.post_transforms_processing(sample)
        return sample

    def prep_for_transforms(self, sample):
        return {"image": np.array(sample)}

    def post_transforms_processing(self, sample):
        return sample

import unittest

import numpy as np

from super_gradients.training.datasets import Cifar10, Cifar100, ImageNetDataset
from albumentations import Compose, HorizontalFlip, InvertImg


class AlbumentationsIntegrationTest(unittest.TestCase):
    def _apply_aug(self, img_no_aug):
        pipe = Compose(transforms=[HorizontalFlip(p=1.0), InvertImg(p=1.0)])
        img_no_aug_transformed = pipe(image=np.array(img_no_aug))["image"]
        return img_no_aug_transformed

    def test_cifar10_albumentations_integration(self):
        ds_no_aug = Cifar10(root="./data/cifar10", train=True, download=True)
        img_no_aug, _ = ds_no_aug.__getitem__(0)

        ds = Cifar10(
            root="./data/cifar10",
            train=True,
            download=True,
            transforms={"Albumentations": {"Compose": {"transforms": [{"HorizontalFlip": {"p": 1.0}}, {"InvertImg": {"p": 1.0}}]}}},
        )

        img_aug, _ = ds.__getitem__(0)
        img_no_aug_transformed = self._apply_aug(img_no_aug)

        self.assertTrue(np.allclose(img_no_aug_transformed, img_aug))

    def test_cifar100_albumentations_integration(self):
        ds_no_aug = Cifar100(root="./data/cifar100", train=True, download=True)
        img_no_aug, _ = ds_no_aug.__getitem__(0)

        ds = Cifar100(
            root="./data/cifar100",
            train=True,
            download=True,
            transforms={"Albumentations": {"Compose": {"transforms": [{"HorizontalFlip": {"p": 1}}, {"InvertImg": {"p": 1.0}}]}}},
        )

        img_aug, _ = ds.__getitem__(0)
        img_no_aug_transformed = self._apply_aug(img_no_aug)

        self.assertTrue(np.allclose(img_no_aug_transformed, img_aug))

    def test_imagenet_albumentations_integration(self):
        ds_no_aug = ImageNetDataset(root="/data/Imagenet/val")
        img_no_aug, _ = ds_no_aug.__getitem__(0)

        ds = ImageNetDataset(
            root="/data/Imagenet/val", transforms={"Albumentations": {"Compose": {"transforms": [{"HorizontalFlip": {"p": 1}}, {"InvertImg": {"p": 1.0}}]}}}
        )
        img_aug, _ = ds.__getitem__(0)
        img_no_aug_transformed = self._apply_aug(img_no_aug)

        self.assertTrue(np.allclose(img_no_aug_transformed, img_aug))


if __name__ == "__main__":
    unittest.main()