[WIP] Mix up (#1549)

* Cleanup * Cleanup * Added MixUp tranform * Update tests/test_transforms.py Co-authored-by: sourcery-ai[bot] <58596630+sourcery-ai[bot]@users.noreply.github.com> * Update albumentations/core/transforms_interface.py Co-authored-by: sourcery-ai[bot] <58596630+sourcery-ai[bot]@users.noreply.github.com> * Update tests/test_transforms.py Co-authored-by: sourcery-ai[bot] <58596630+sourcery-ai[bot]@users.noreply.github.com> * Fix in docstring * Fix in docstring * Fixes from Code Review * Works with iterator * Added test for multiple targets * Added test for multiple targets * Added test for multiple targets * Added test for multiple targets * Fix in docs * Fix in CI * Refactoring image generation * Updated way we find targets of transforms * Added test for targets in docstrings * Added Mixing transforms to docs * Cleanup * removed if_pytorch from conftest --------- Co-authored-by: sourcery-ai[bot] <58596630+sourcery-ai[bot]@users.noreply.github.com>
albumentations-team · Mar 3, 2024 · 138bc5c · 138bc5c
1 parent 9b2863b
commit 138bc5c
Show file tree

Hide file tree

Showing 34 changed files with 1,039 additions and 226 deletions.
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
@@ -41,7 +41,7 @@ jobs:
       run: pip install -r requirements-dev.txt
     - name: Install dependencies
       run: |
-        pip install .[tests]
+        pip install .
     - name: Cleanup the build directory
       uses: JesseTG/[email protected]
       with:
@@ -88,6 +88,6 @@ jobs:
         pip install requests
         pip install .
     - name: Run checks for documentation
-      run: python tools/make_transforms_docs.py check README.md
+      run: python -m tools.make_transforms_docs check README.md
     - name: Run checks for used by documentation
-      run: python tools/make_used_by_docs.py check
+      run: python -m tools.make_used_by_docs check
diff --git a/README.md b/README.md
diff --git a/albumentations/augmentations/__init__.py b/albumentations/augmentations/__init__.py
@@ -18,5 +18,7 @@
 from .geometric.resize import *
 from .geometric.rotate import *
 from .geometric.transforms import *
+from .mixing.functional import *
+from .mixing.transforms import *
 from .transforms import *
 from .utils import *
diff --git a/albumentations/augmentations/crops/transforms.py b/albumentations/augmentations/crops/transforms.py
@@ -8,7 +8,7 @@
 from albumentations.augmentations.geometric import functional as FGeometric
 from albumentations.core.bbox_utils import union_of_bboxes
 from albumentations.core.transforms_interface import DualTransform, to_tuple
-from albumentations.core.types import BoxInternalType, KeypointInternalType, ScaleFloatType
+from albumentations.core.types import BoxInternalType, KeypointInternalType, ScaleFloatType, Targets
 
 from . import functional as F
 
@@ -47,6 +47,8 @@ class RandomCrop(DualTransform):
 
     """
 
+    _targets = (Targets.IMAGE, Targets.MASK, Targets.BBOXES, Targets.KEYPOINTS)
+
     def __init__(self, height: int, width: int, always_apply: bool = False, p: float = 1.0):
         super().__init__(always_apply, p)
         self.height = height
@@ -91,6 +93,8 @@ class CenterCrop(DualTransform):
 
     """
 
+    _targets = (Targets.IMAGE, Targets.MASK, Targets.BBOXES, Targets.KEYPOINTS)
+
     def __init__(self, height: int, width: int, always_apply: bool = False, p: float = 1.0):
         super().__init__(always_apply, p)
         self.height = height
@@ -127,6 +131,8 @@ class Crop(DualTransform):
 
     """
 
+    _targets = (Targets.IMAGE, Targets.MASK, Targets.BBOXES, Targets.KEYPOINTS)
+
     def __init__(
         self,
         x_min: int = 0,
@@ -176,6 +182,8 @@ class CropNonEmptyMaskIfExists(DualTransform):
 
     """
 
+    _targets = (Targets.IMAGE, Targets.MASK, Targets.BBOXES, Targets.KEYPOINTS)
+
     def __init__(
         self,
         height: int,
@@ -351,6 +359,8 @@ class RandomSizedCrop(_BaseRandomSizedCrop):
 
     """
 
+    _targets = (Targets.IMAGE, Targets.MASK, Targets.BBOXES, Targets.KEYPOINTS)
+
     def __init__(
         self,
         min_max_height: Tuple[int, int],
@@ -362,6 +372,7 @@ def __init__(
         p: float = 1.0,
     ):
         super().__init__(height=height, width=width, interpolation=interpolation, always_apply=always_apply, p=p)
+
         self.min_max_height = min_max_height
         self.w2h_ratio = w2h_ratio
 
@@ -400,6 +411,8 @@ class RandomResizedCrop(_BaseRandomSizedCrop):
 
     """
 
+    _targets = (Targets.IMAGE, Targets.MASK, Targets.BBOXES, Targets.KEYPOINTS)
+
     def __init__(
         self,
         height: int,
@@ -493,6 +506,8 @@ class RandomCropNearBBox(DualTransform):
 
     """
 
+    _targets = (Targets.IMAGE, Targets.MASK, Targets.BBOXES, Targets.KEYPOINTS)
+
     def __init__(
         self,
         max_part_shift: ScaleFloatType = (0.3, 0.3),
@@ -564,6 +579,8 @@ class BBoxSafeRandomCrop(DualTransform):
 
     """
 
+    _targets = (Targets.IMAGE, Targets.MASK, Targets.BBOXES)
+
     def __init__(self, erosion_rate: float = 0.0, always_apply: bool = False, p: float = 1.0):
         super().__init__(always_apply, p)
         self.erosion_rate = erosion_rate
@@ -644,6 +661,8 @@ class RandomSizedBBoxSafeCrop(BBoxSafeRandomCrop):
 
     """
 
+    _targets = (Targets.IMAGE, Targets.MASK, Targets.BBOXES)
+
     def __init__(
         self,
         height: int,
@@ -765,6 +784,8 @@ class CropAndPad(DualTransform):
 
     """
 
+    _targets = (Targets.IMAGE, Targets.MASK, Targets.BBOXES, Targets.KEYPOINTS)
+
     def __init__(
         self,
         px: Optional[Union[int, List[int]]] = None,
@@ -1029,6 +1050,8 @@ class RandomCropFromBorders(DualTransform):
 
     """
 
+    _targets = (Targets.IMAGE, Targets.MASK, Targets.BBOXES, Targets.KEYPOINTS)
+
     def __init__(
         self,
         crop_left: float = 0.1,

diff --git a/albumentations/augmentations/dropout/coarse_dropout.py b/albumentations/augmentations/dropout/coarse_dropout.py
@@ -4,7 +4,7 @@
 import numpy as np
 
 from albumentations.core.transforms_interface import DualTransform
-from albumentations.core.types import KeypointType, ScalarType
+from albumentations.core.types import KeypointType, ScalarType, Targets
 
 from .functional import cutout, keypoint_in_hole
 
@@ -47,6 +47,8 @@ class CoarseDropout(DualTransform):
 
     """
 
+    _targets = (Targets.IMAGE, Targets.MASK, Targets.KEYPOINTS)
+
     def __init__(
         self,
         max_holes: int = 8,

diff --git a/albumentations/augmentations/dropout/grid_dropout.py b/albumentations/augmentations/dropout/grid_dropout.py
@@ -4,7 +4,7 @@
 import numpy as np
 
 from albumentations.core.transforms_interface import DualTransform
-from albumentations.core.types import ScalarType
+from albumentations.core.types import ScalarType, Targets
 
 from . import functional as F
 
@@ -50,6 +50,8 @@ class GridDropout(DualTransform):
 
     """
 
+    _targets = (Targets.IMAGE, Targets.MASK)
+
     def __init__(
         self,
         ratio: float = 0.5,

diff --git a/albumentations/augmentations/dropout/mask_dropout.py b/albumentations/augmentations/dropout/mask_dropout.py
@@ -6,7 +6,7 @@
 from skimage.measure import label
 
 from albumentations.core.transforms_interface import DualTransform, to_tuple
-from albumentations.core.types import ScalarType
+from albumentations.core.types import ScalarType, Targets
 
 __all__ = ["MaskDropout"]
 
@@ -24,7 +24,7 @@ class MaskDropout(DualTransform):
     ----
         max_objects: Maximum number of labels that can be zeroed out. Can be tuple, in this case it's [min, max]
         image_fill_value: Fill value to use when filling image.
-            Can be 'inpaint' to apply inpaining (works only  for 3-chahnel images)
+            Can be 'inpaint' to apply inpainting (works only  for 3-channel images)
         mask_fill_value: Fill value to use when filling mask.
 
     Targets:
@@ -35,6 +35,8 @@ class MaskDropout(DualTransform):
 
     """
 
+    _targets = (Targets.IMAGE, Targets.MASK)
+
     def __init__(
         self,
         max_objects: int = 1,

diff --git a/albumentations/augmentations/dropout/xy_masking.py b/albumentations/augmentations/dropout/xy_masking.py
@@ -4,7 +4,7 @@
 import numpy as np
 
 from albumentations.core.transforms_interface import DualTransform
-from albumentations.core.types import ColorType, KeypointType, ScaleIntType
+from albumentations.core.types import ColorType, KeypointType, ScaleIntType, Targets
 
 from .functional import cutout, keypoint_in_hole
 
@@ -49,6 +49,8 @@ class XYMasking(DualTransform):
 
     """
 
+    _targets = (Targets.IMAGE, Targets.MASK, Targets.KEYPOINTS)
+
     def __init__(
         self,
         num_masks_x: ScaleIntType = 0,

diff --git a/albumentations/augmentations/geometric/functional.py b/albumentations/augmentations/geometric/functional.py
@@ -15,8 +15,7 @@
     preserve_shape,
 )
 from albumentations.core.bbox_utils import denormalize_bbox, normalize_bbox
-from albumentations.core.transforms_interface import FillValueType
-from albumentations.core.types import BoxInternalType, ImageColorType, KeypointInternalType
+from albumentations.core.types import BoxInternalType, ColorType, ImageColorType, KeypointInternalType
 
 __all__ = [
     "optical_distortion",
@@ -654,14 +653,14 @@ def safe_rotate(
     img: np.ndarray,
     matrix: np.ndarray,
     interpolation: int,
-    value: FillValueType = None,
+    value: Optional[ColorType] = None,
     border_mode: int = cv2.BORDER_REFLECT_101,
 ) -> np.ndarray:
-    h, w = img.shape[:2]
+    height, width = img.shape[:2]
     warp_fn = _maybe_process_in_chunks(
         cv2.warpAffine,
         M=matrix,
-        dsize=(w, h),
+        dsize=(width, height),
         flags=interpolation,
         borderMode=border_mode,
         borderValue=value,

diff --git a/albumentations/augmentations/geometric/resize.py b/albumentations/augmentations/geometric/resize.py
@@ -5,7 +5,7 @@
 import numpy as np
 
 from albumentations.core.transforms_interface import DualTransform, to_tuple
-from albumentations.core.types import BoxInternalType, KeypointInternalType, ScaleFloatType
+from albumentations.core.types import BoxInternalType, KeypointInternalType, ScaleFloatType, Targets
 
 from . import functional as F
 
@@ -34,6 +34,8 @@ class RandomScale(DualTransform):
 
     """
 
+    _targets = (Targets.IMAGE, Targets.MASK, Targets.BBOXES, Targets.KEYPOINTS)
+
     def __init__(
         self,
         scale_limit: ScaleFloatType = 0.1,
@@ -84,6 +86,8 @@ class LongestMaxSize(DualTransform):
 
     """
 
+    _targets = (Targets.IMAGE, Targets.MASK, Targets.BBOXES, Targets.KEYPOINTS)
+
     def __init__(
         self,
         max_size: Union[int, Sequence[int]] = 1024,
@@ -138,6 +142,8 @@ class SmallestMaxSize(DualTransform):
 
     """
 
+    _targets = (Targets.IMAGE, Targets.MASK, Targets.KEYPOINTS, Targets.BBOXES)
+
     def __init__(
         self,
         max_size: Union[int, Sequence[int]] = 1024,
@@ -193,6 +199,8 @@ class Resize(DualTransform):
 
     """
 
+    _targets = (Targets.IMAGE, Targets.MASK, Targets.KEYPOINTS, Targets.BBOXES)
+
     def __init__(
         self, height: int, width: int, interpolation: int = cv2.INTER_LINEAR, always_apply: bool = False, p: float = 1
     ):

diff --git a/albumentations/augmentations/geometric/rotate.py b/albumentations/augmentations/geometric/rotate.py
@@ -1,13 +1,13 @@
 import math
 import random
-from typing import Any, Dict, List, Optional, Sequence, Tuple, Union, cast
+from typing import Any, Dict, List, Optional, Tuple, Union, cast
 
 import cv2
 import numpy as np
 
 from albumentations.augmentations.crops import functional as FCrops
-from albumentations.core.transforms_interface import DualTransform, FillValueType, to_tuple
-from albumentations.core.types import BoxInternalType, KeypointInternalType, ScaleIntType
+from albumentations.core.transforms_interface import DualTransform, to_tuple
+from albumentations.core.types import BoxInternalType, ColorType, KeypointInternalType, ScaleIntType, Targets
 
 from . import functional as F
 
@@ -31,6 +31,8 @@ class RandomRotate90(DualTransform):
 
     """
 
+    _targets = (Targets.IMAGE, Targets.MASK, Targets.BBOXES, Targets.KEYPOINTS)
+
     def apply(self, img: np.ndarray, factor: float = 0, **params: Any) -> np.ndarray:
         """Args:
         ----
@@ -83,6 +85,8 @@ class Rotate(DualTransform):
 
     """
 
+    _targets = (Targets.IMAGE, Targets.MASK, Targets.BBOXES, Targets.KEYPOINTS)
+
     def __init__(
         self,
         limit: ScaleIntType = 90,
@@ -250,13 +254,15 @@ class SafeRotate(DualTransform):
 
     """
 
+    _targets = (Targets.IMAGE, Targets.MASK, Targets.BBOXES, Targets.KEYPOINTS)
+
     def __init__(
         self,
         limit: Union[float, Tuple[float, float]] = 90,
         interpolation: int = cv2.INTER_LINEAR,
         border_mode: int = cv2.BORDER_REFLECT_101,
-        value: FillValueType = None,
-        mask_value: Optional[Union[int, float, Sequence[int], Sequence[float]]] = None,
+        value: Optional[ColorType] = None,
+        mask_value: Optional[ColorType] = None,
         always_apply: bool = False,
         p: float = 0.5,
     ):
@@ -296,10 +302,10 @@ def get_params_dependent_on_targets(self, params: Dict[str, Any]) -> Dict[str, A
         angle = random.uniform(self.limit[0], self.limit[1])
 
         image = params["image"]
-        h, w = image.shape[:2]
+        height, width = image.shape[:2]
 
         # https://stackoverflow.com/questions/43892506/opencv-python-rotate-image-without-cropping-sides
-        image_center = (w / 2, h / 2)
+        image_center = (width / 2, height / 2)
 
         # Rotation Matrix
         rotation_mat = cv2.getRotationMatrix2D(image_center, angle, 1.0)
@@ -309,11 +315,11 @@ def get_params_dependent_on_targets(self, params: Dict[str, Any]) -> Dict[str, A
         abs_sin = abs(rotation_mat[0, 1])
 
         # find the new width and height bounds
-        new_w = math.ceil(h * abs_sin + w * abs_cos)
-        new_h = math.ceil(h * abs_cos + w * abs_sin)
+        new_w = math.ceil(height * abs_sin + width * abs_cos)
+        new_h = math.ceil(height * abs_cos + width * abs_sin)
 
-        scale_x = w / new_w
-        scale_y = h / new_h
+        scale_x = width / new_w
+        scale_y = height / new_h
 
         # Shift the image to create padding
         rotation_mat[0, 2] += new_w / 2 - image_center[0]