Add albumentations to registry (#1238)

* added albumentations to registry * removed unnecessary print * codestyle * removed omegaconf * updated requirements * updated requirements [2] * updated requirements [3] * added 'cache pip contrib [2]' to github workflow * added version for albumentations * added version for albumentations; updated CONTRIBUTING.md * update .github/workflows * fixed line length for codestyle
catalyst-team · Jun 24, 2021 · bc7ae8b · bc7ae8b
1 parent 7f63545
commit bc7ae8b
Show file tree

Hide file tree

Showing 14 changed files with 79 additions and 10 deletions.
diff --git a/.github/workflows/codestyle.yml b/.github/workflows/codestyle.yml
@@ -57,10 +57,18 @@ jobs:
           restore-keys: |
             ${{ runner.os }}-${{ matrix.python-version }}-pip-contrib
 
+      - name: cache pip contrib [2]
+        uses: actions/cache@v2
+        with:
+          path: ${{ steps.pip-cache.outputs.dir }}
+          key: ${{ runner.os }}-${{ matrix.python-version }}-pip-contrib -${{ hashFiles('./requirements/requirements-albu.txt') }}
+          restore-keys: |
+            ${{ runner.os }}-${{ matrix.python-version }}-pip-contrib
+
       - name: install dependencies
         run: |
           # python -m pip install --upgrade --user pip
-          pip install -r ./requirements/requirements.txt -r ./requirements/requirements-cv.txt -r ./requirements/requirements-dev.txt -r ./requirements/requirements-hydra.txt -r ./requirements/requirements-ml.txt -r ./requirements/requirements-nifti.txt -r ./requirements/requirements-optuna.txt -r ./requirements/requirements-mlflow.txt -r ./requirements/requirements-neptune.txt
+          pip install -r ./requirements/requirements.txt -r ./requirements/requirements-cv.txt -r ./requirements/requirements-dev.txt -r ./requirements/requirements-hydra.txt -r ./requirements/requirements-ml.txt -r ./requirements/requirements-nifti.txt -r ./requirements/requirements-optuna.txt -r ./requirements/requirements-mlflow.txt -r ./requirements/requirements-neptune.txt -r ./requirements/requirements-albu.txt
           python --version
           pip --version
           pip list
@@ -119,10 +127,18 @@ jobs:
           restore-keys: |
             ${{ runner.os }}-${{ matrix.python-version }}-pip-contrib
 
+      - name: cache pip contrib [2]
+        uses: actions/cache@v2
+        with:
+          path: ${{ steps.pip-cache.outputs.dir }}
+          key: ${{ runner.os }}-${{ matrix.python-version }}-pip-contrib -${{ hashFiles('./requirements/requirements-albu.txt') }}
+          restore-keys: |
+            ${{ runner.os }}-${{ matrix.python-version }}-pip-contrib
+
       - name: install dependencies
         run: |
           # python -m pip install --upgrade --user pip
-          pip install -r ./requirements/requirements.txt -r ./requirements/requirements-cv.txt -r ./requirements/requirements-dev.txt -r ./requirements/requirements-hydra.txt -r ./requirements/requirements-ml.txt -r ./requirements/requirements-optuna.txt -r ./requirements/requirements-mlflow.txt -r ./requirements/requirements-neptune.txt
+          pip install -r ./requirements/requirements.txt -r ./requirements/requirements-cv.txt -r ./requirements/requirements-dev.txt -r ./requirements/requirements-hydra.txt -r ./requirements/requirements-ml.txt -r ./requirements/requirements-optuna.txt -r ./requirements/requirements-mlflow.txt -r ./requirements/requirements-neptune.txt  -r ./requirements/requirements-albu.txt
           python --version
           pip --version
           pip list

diff --git a/.github/workflows/deploy_push.yml b/.github/workflows/deploy_push.yml
@@ -55,14 +55,14 @@ jobs:
         uses: actions/cache@v2
         with:
           path: ${{ steps.pip-cache.outputs.dir }}
-          key: ${{ runner.os }}-${{ matrix.python-version }}-pip -${{hashFiles('./requirements/requirements-cv.txt') }} -${{hashFiles('./requirements/requirements-nifti.txt') }} -${{ hashFiles('./requirements/requirements-hydra.txt') }} -${{ hashFiles('./requirements/requirements-ml.txt') }} -${{ hashFiles('./requirements/requirements-optuna.txt') }} -${{ hashFiles('./requirements/requirements-mlflow.txt') }} -${{ hashFiles('./requirements/requirements-neptune.txt') }}
+          key: ${{ runner.os }}-${{ matrix.python-version }}-pip -${{hashFiles('./requirements/requirements-cv.txt') }} -${{hashFiles('./requirements/requirements-nifti.txt') }} -${{ hashFiles('./requirements/requirements-hydra.txt') }} -${{ hashFiles('./requirements/requirements-ml.txt') }} -${{ hashFiles('./requirements/requirements-optuna.txt') }} -${{ hashFiles('./requirements/requirements-mlflow.txt') }} -${{ hashFiles('./requirements/requirements-neptune.txt') }} -${{ hashFiles('./requirements/requirements-albu.txt') }}
           restore-keys: |
             ${{ runner.os }}-${{ matrix.python-version }}-pip-contrib
 
       - name: install dependencies
         run: |
           # python -m pip install --upgrade --user pip
-          pip install -r ./requirements/requirements.txt -r ./requirements/requirements-cv.txt -r ./requirements/requirements-dev.txt -r ./requirements/requirements-hydra.txt -r ./requirements/requirements-ml.txt -r ./requirements/requirements-nifti.txt -r ./requirements/requirements-optuna.txt -r ./requirements/requirements-mlflow.txt -r ./requirements/requirements-neptune.txt
+          pip install -r ./requirements/requirements.txt -r ./requirements/requirements-cv.txt -r ./requirements/requirements-dev.txt -r ./requirements/requirements-hydra.txt -r ./requirements/requirements-ml.txt -r ./requirements/requirements-nifti.txt -r ./requirements/requirements-optuna.txt -r ./requirements/requirements-mlflow.txt -r ./requirements/requirements-neptune.txt -r ./requirements/requirements-albu.txt
           python --version
           pip --version
           pip list

diff --git a/.github/workflows/dl_cpu.yml b/.github/workflows/dl_cpu.yml
@@ -68,6 +68,7 @@ jobs:
           python -c "req = open('./requirements/requirements-cv.txt').read().replace('>', '=') ; open('./requirements/requirements-cv.txt', 'w').write(req)"
           python -c "req = open('./requirements/requirements-ml.txt').read().replace('>', '=') ; open('./requirements/requirements-ml.txt', 'w').write(req)"
           python -c "req = open('./requirements/requirements-neptune.txt').read().replace('>', '=') ; open('./requirements/requirements-neptune.txt', 'w').write(req)"
+          python -c "req = open('./requirements/requirements-albu.txt').read().replace('>', '=') ; open('./requirements/requirements-albu.txt', 'w').write(req)"
 
       # https://github.com/actions/cache/blob/master/examples.md
       # Note: This uses an internal pip API and may not always work
@@ -96,7 +97,7 @@ jobs:
       - name: install dependencies
         run: |
           # python -m pip install --upgrade --user pip
-          pip install -r ./requirements/requirements.txt -r ./requirements/requirements-cv.txt -r ./requirements/requirements-dev.txt -r ./requirements/requirements-hydra.txt -r ./requirements/requirements-ml.txt -r ./requirements/requirements-optuna.txt -r ./requirements/requirements-mlflow.txt -r ./requirements/requirements-nifti.txt -r ./requirements/requirements-neptune.txt
+          pip install -r ./requirements/requirements.txt -r ./requirements/requirements-cv.txt -r ./requirements/requirements-dev.txt -r ./requirements/requirements-hydra.txt -r ./requirements/requirements-ml.txt -r ./requirements/requirements-optuna.txt -r ./requirements/requirements-mlflow.txt -r ./requirements/requirements-nifti.txt -r ./requirements/requirements-neptune.txt -r ./requirements/requirements-albu.txt
           python --version
           pip --version
           pip list

diff --git a/.github/workflows/dl_cpu_extras.yml b/.github/workflows/dl_cpu_extras.yml
@@ -22,7 +22,7 @@ jobs:
         os: [ubuntu-20.04]
         python-version: [3.7]
         requirements: ['latest']
-        extras: ['base', 'cv', 'ml', 'hydra', 'optuna', 'onnx', 'neptune']
+        extras: ['base', 'cv', 'ml', 'hydra', 'optuna', 'onnx', 'neptune', 'albu']
 
     timeout-minutes: 30
     steps:
@@ -103,6 +103,16 @@ jobs:
           pip list
         shell: bash
 
+      - name: set albu extras
+        if: matrix.extras == 'albu'
+        run: |
+          # python -m pip install --upgrade --user pip
+          pip install -r ./requirements/requirements.txt -r ./requirements/requirements-dev.txt -r ./requirements/requirements-albu.txt
+          python --version
+          pip --version
+          pip list
+        shell: bash
+
 # <- standard block end ->
 
       - name: check examples

diff --git a/.github/workflows/dl_cpu_minimal.yml b/.github/workflows/dl_cpu_minimal.yml
@@ -68,6 +68,7 @@ jobs:
           python -c "req = open('./requirements/requirements-cv.txt').read().replace('>', '=') ; open('./requirements/requirements-cv.txt', 'w').write(req)"
           python -c "req = open('./requirements/requirements-ml.txt').read().replace('>', '=') ; open('./requirements/requirements-ml.txt', 'w').write(req)"
           python -c "req = open('./requirements/requirements-neptune.txt').read().replace('>', '=') ; open('./requirements/requirements-neptune.txt', 'w').write(req)"
+           python -c "req = open('./requirements/requirements-albu.txt').read().replace('>', '=') ; open('./requirements/requirements-albu.txt', 'w').write(req)"
 
       # https://github.com/actions/cache/blob/master/examples.md
       # Note: This uses an internal pip API and may not always work
@@ -96,7 +97,7 @@ jobs:
       - name: install dependencies
         run: |
           # python -m pip install --upgrade --user pip
-          pip install -r ./requirements/requirements.txt -r ./requirements/requirements-cv.txt -r ./requirements/requirements-dev.txt -r ./requirements/requirements-hydra.txt -r ./requirements/requirements-ml.txt -r ./requirements/requirements-optuna.txt -r ./requirements/requirements-mlflow.txt -r ./requirements/requirements-nifti.txt -r ./requirements/requirements-neptune.txt
+          pip install -r ./requirements/requirements.txt -r ./requirements/requirements-cv.txt -r ./requirements/requirements-dev.txt -r ./requirements/requirements-hydra.txt -r ./requirements/requirements-ml.txt -r ./requirements/requirements-optuna.txt -r ./requirements/requirements-mlflow.txt -r ./requirements/requirements-nifti.txt -r ./requirements/requirements-neptune.txt -r ./requirements/requirements-albu.txt
           python --version
           pip --version
           pip list

diff --git a/.github/workflows/integrations.yml b/.github/workflows/integrations.yml
@@ -60,6 +60,7 @@ jobs:
           python -c "req = open('./requirements/requirements-optuna.txt').read().replace('>', '=') ; open('./requirements/requirements-optuna.txt', 'w').write(req)"
           python -c "req = open('./requirements/requirements-nifti.txt').read().replace('>', '=') ; open('./requirements/requirements-nifti.txt', 'w').write(req)"
           python -c "req = open('./requirements/requirements-neptune.txt').read().replace('>', '=') ; open('./requirements/requirements-neptune.txt', 'w').write(req)"
+          python -c "req = open('./requirements/requirements-albu.txt').read().replace('>', '=') ; open('./requirements/requirements-albu.txt', 'w').write(req)"
 
       # https://github.com/actions/cache/blob/master/examples.md
       # Note: This uses an internal pip API and may not always work
@@ -85,10 +86,18 @@ jobs:
           restore-keys: |
             ${{ runner.os }}-${{ matrix.python-version }}-pip-contrib
 
+      - name: cache pip contrib [2]
+        uses: actions/cache@v2
+        with:
+          path: ${{ steps.pip-cache.outputs.dir }}
+          key: ${{ runner.os }}-${{ matrix.python-version }}-pip -${{ hashFiles('./requirements/requirements-albu.txt') }}
+          restore-keys: |
+            ${{ runner.os }}-${{ matrix.python-version }}-pip-contrib
+
       - name: install dependencies
         run: |
           # python -m pip install --upgrade --user pip
-          pip install -r ./requirements/requirements.txt -r ./requirements/requirements-cv.txt -r ./requirements/requirements-dev.txt -r ./requirements/requirements-hydra.txt -r ./requirements/requirements-ml.txt -r ./requirements/requirements-optuna.txt -r ./requirements/requirements-mlflow.txt -r ./requirements/requirements-neptune.txt
+          pip install -r ./requirements/requirements.txt -r ./requirements/requirements-cv.txt -r ./requirements/requirements-dev.txt -r ./requirements/requirements-hydra.txt -r ./requirements/requirements-ml.txt -r ./requirements/requirements-optuna.txt -r ./requirements/requirements-mlflow.txt -r ./requirements/requirements-neptune.txt  -r ./requirements/requirements-albu.txt
           python --version
           pip --version
           pip list

diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
@@ -88,7 +88,8 @@ If you are not familiar with creating a Pull Request, here are some guides:
       -r requirements/requirements-nifti.txt \
       -r requirements/requirements-onnx.txt \
       -r requirements/requirements-optuna.txt \
-      -r requirements/requirements-neptune.txt
+      -r requirements/requirements-neptune.txt \
+      -r requirements/requirements-albu.txt
     ```
 2. Break your work into small, single-purpose updates if possible.
 It's much harder to merge in a large change with a lot of disjoint features.

diff --git a/bin/workflows/check_settings.sh b/bin/workflows/check_settings.sh
@@ -11,6 +11,7 @@ function clean_requirements() {
   pip uninstall -r requirements/requirements-optuna.txt -y
   pip uninstall -r requirements/requirements-mlflow.txt -y
   pip uninstall -r requirements/requirements-neptune.txt -y
+  pip uninstall -r requirements/requirements-albu.txt -y
   pip install -r requirements/requirements.txt --quiet \
   --find-links https://download.pytorch.org/whl/cpu/torch_stable.html \
   --upgrade-strategy only-if-needed
@@ -28,6 +29,7 @@ hydra_required = false
 optuna_required = false
 mlflow_required = false
 neptune_required = false
+albu_required = false
 ")"
   for REQUIRED in "$@"
   do
@@ -48,6 +50,7 @@ ml_required = false
 hydra_required = false
 optuna_required = false
 mlflow_required = false
+albu_required = false
 EOT
 
 python -c """
@@ -172,6 +175,7 @@ ml_required = true
 hydra_required = false
 mlflow_required = false
 neptune_required = false
+albu_required = false
 EOT
 
 python -c """

diff --git a/catalyst/registry.py b/catalyst/registry.py
@@ -14,6 +14,15 @@ def _transforms_loader(r: registry.Registry):
 
     r.add_from_module(t, prefix=["catalyst.", "C."])
 
+    if SETTINGS.albu_required:
+        import albumentations as m
+
+        r.add_from_module(m, prefix=["A.", "albu.", "albumentations."])
+
+        from albumentations import pytorch as p
+
+        r.add_from_module(p, prefix=["A.", "albu.", "albumentations."])
+
 
 REGISTRY.late_add(_transforms_loader)
 

diff --git a/catalyst/settings.py b/catalyst/settings.py
@@ -33,6 +33,15 @@ def _is_amp_available():
         return False
 
 
+def _is_albumentations_available():
+    try:
+        import albumentations as albu  # noqa: F401
+
+        return True
+    except ModuleNotFoundError:
+        return False
+
+
 def _is_xla_available():
     try:
         import torch_xla.core.xla_model as xm  # noqa: F401
@@ -189,6 +198,7 @@ def __init__(  # noqa: D107
         nifti_required: Optional[bool] = None,
         ml_required: Optional[bool] = None,
         # [integrations]
+        albu_required: Optional[bool] = None,
         hydra_required: Optional[bool] = None,
         # nmslib_required: Optional[bool] = False,
         optuna_required: Optional[bool] = None,
@@ -234,6 +244,11 @@ def __init__(  # noqa: D107
         )
 
         # [integrations]
+        self.albu_required: bool = _get_optional_value(
+            albu_required,
+            _is_albumentations_available,
+            "catalyst[albu] is not available, to install it, " "run `pip install catalyst[albu]`.",
+        )
         self.hydra_required: bool = _get_optional_value(
             hydra_required,
             _is_hydra_available,

diff --git a/docker/Dockerfile b/docker/Dockerfile
@@ -44,6 +44,7 @@ RUN pip install -U catalyst --no-cache-dir
 RUN if [ "$CATALYST_DEV" = "1" ] ; then pip install -U catalyst[dev] --no-cache-dir ; else echo "catalyst[dev] is not required" ; fi
 RUN if [ "$CATALYST_CV" = "1" ] ; then pip install -U catalyst[cv] --no-cache-dir ; else echo "catalyst[cv] is not required" ; fi
 RUN if [ "$CATALYST_ML" = "1" ] ; then pip install -U catalyst[ml] --no-cache-dir ; else echo "catalyst[ml] is not required" ; fi
+RUN if [ "$CATALYST_ALBU" = "1" ] ; then pip install -U catalyst[albu] --no-cache-dir ; else echo "catalyst[albu] is not required" ; fi
 RUN if [ "$CATALYST_HYDRA" = "1" ] ; then pip install -U catalyst[hydra] --no-cache-dir ; else echo "catalyst[hydra] is not required" ; fi
 RUN if [ "$CATALYST_OPTUNA" = "1" ] ; then pip install -U catalyst[optuna] --no-cache-dir ; else echo "catalyst[optuna] is not required" ; fi
 RUN if [ "$CATALYST_ONNX" = "1" ] ; then pip install -U catalyst[onnx] --no-cache-dir ; else echo "catalyst[onnx] is not required" ; fi

diff --git a/requirements/requirements-albu.txt b/requirements/requirements-albu.txt
@@ -0,0 +1 @@
+albumentations>=0.5.1
diff --git a/requirements/requirements-cv.txt b/requirements/requirements-cv.txt
@@ -1,5 +1,5 @@
 imageio>=2.5.0
-opencv-python-headless>=3.4.5.20
+opencv-python-headless>=4.1.1.26
 scikit-image>=0.16.1
 torchvision>=0.4.1
 Pillow>=6.1  # torchvision fix (https://github.com/python-pillow/Pillow/issues/4130)
diff --git a/setup.py b/setup.py
@@ -99,6 +99,7 @@ def run(self):
     "wandb": load_requirements("requirements/requirements-wandb.txt"),
     "fairscale": load_requirements("requirements/requirements-fairscale.txt"),
     "deepspeed": load_requirements("requirements/requirements-deepspeed.txt"),
+    "albu": load_requirements("requirements/requirements-albu.txt"),
 }
 extras["all"] = extras["cv"] + extras["ml"] + extras["hydra"] + extras["optuna"]
 # Meta dependency groups.