add BBoxSafeRandomCrop (#579)

* add BBoxSafeRandomCrop w/o resizing * add new transform to `__all__` * update new transform in README * update docstring of BBoxSafeRandomCrop * Fix tests Co-authored-by: Dipet <dipetm@gmail.com>
albumentations-team · Aug 16, 2022 · a271a09 · a271a09
1 parent fe4adaf
commit a271a09
Show file tree

Hide file tree

Showing 5 changed files with 64 additions and 25 deletions.
diff --git a/README.md b/README.md
@@ -173,6 +173,7 @@ Spatial-level transforms will simultaneously change both an input image as well
 | Transform                                                                                                                                                                       | Image | Masks | BBoxes | Keypoints |
 | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | :---: | :---: | :----: | :-------: |
 | [Affine](https://albumentations.ai/docs/api_reference/augmentations/geometric/transforms/#albumentations.augmentations.geometric.transforms.Affine)                             | ✓     | ✓     | ✓      | ✓         |
+| [BBoxSafeRandomCrop](https://albumentations.ai/docs/api_reference/augmentations/crops/transforms/#albumentations.augmentations.crops.transforms.BBoxSafeRandomCrop)             | ✓     | ✓     | ✓      |           |
 | [CenterCrop](https://albumentations.ai/docs/api_reference/augmentations/crops/transforms/#albumentations.augmentations.crops.transforms.CenterCrop)                             | ✓     | ✓     | ✓      | ✓         |
 | [CoarseDropout](https://albumentations.ai/docs/api_reference/augmentations/dropout/coarse_dropout/#albumentations.augmentations.dropout.coarse_dropout.CoarseDropout)           | ✓     | ✓     |        | ✓         |
 | [Crop](https://albumentations.ai/docs/api_reference/augmentations/crops/transforms/#albumentations.augmentations.crops.transforms.Crop)                                         | ✓     | ✓     | ✓      | ✓         |

diff --git a/albumentations/augmentations/crops/transforms.py b/albumentations/augmentations/crops/transforms.py
@@ -22,6 +22,7 @@
     "RandomSizedBBoxSafeCrop",
     "CropAndPad",
     "RandomCropFromBorders",
+    "BBoxSafeRandomCrop",
 ]
 
 
@@ -471,35 +472,23 @@ def get_transform_init_args_names(self) -> Tuple[str]:
         return ("max_part_shift",)
 
 
-class RandomSizedBBoxSafeCrop(DualTransform):
-    """Crop a random part of the input and rescale it to some size without loss of bboxes.
-
+class BBoxSafeRandomCrop(DualTransform):
+    """Crop a random part of the input without loss of bboxes.
     Args:
-        height (int): height after crop and resize.
-        width (int): width after crop and resize.
         erosion_rate (float): erosion rate applied on input image height before crop.
-        interpolation (OpenCV flag): flag that is used to specify the interpolation algorithm. Should be one of:
-            cv2.INTER_NEAREST, cv2.INTER_LINEAR, cv2.INTER_CUBIC, cv2.INTER_AREA, cv2.INTER_LANCZOS4.
-            Default: cv2.INTER_LINEAR.
         p (float): probability of applying the transform. Default: 1.
-
     Targets:
         image, mask, bboxes
-
     Image types:
         uint8, float32
     """
 
-    def __init__(self, height, width, erosion_rate=0.0, interpolation=cv2.INTER_LINEAR, always_apply=False, p=1.0):
-        super(RandomSizedBBoxSafeCrop, self).__init__(always_apply, p)
-        self.height = height
-        self.width = width
-        self.interpolation = interpolation
+    def __init__(self, erosion_rate=0.0, always_apply=False, p=1.0):
+        super(BBoxSafeRandomCrop, self).__init__(always_apply, p)
         self.erosion_rate = erosion_rate
 
-    def apply(self, img, crop_height=0, crop_width=0, h_start=0, w_start=0, interpolation=cv2.INTER_LINEAR, **params):
-        crop = F.random_crop(img, crop_height, crop_width, h_start, w_start)
-        return FGeometric.resize(crop, self.height, self.width, interpolation)
+    def apply(self, img, crop_height=0, crop_width=0, h_start=0, w_start=0, **params):
+        return F.random_crop(img, crop_height, crop_width, h_start, w_start)
 
     def get_params_dependent_on_targets(self, params):
         img_h, img_w = params["image"].shape[:2]
@@ -534,7 +523,37 @@ def targets_as_params(self):
         return ["image", "bboxes"]
 
     def get_transform_init_args_names(self):
-        return ("height", "width", "erosion_rate", "interpolation")
+        return ("erosion_rate",)
+
+
+class RandomSizedBBoxSafeCrop(BBoxSafeRandomCrop):
+    """Crop a random part of the input and rescale it to some size without loss of bboxes.
+    Args:
+        height (int): height after crop and resize.
+        width (int): width after crop and resize.
+        erosion_rate (float): erosion rate applied on input image height before crop.
+        interpolation (OpenCV flag): flag that is used to specify the interpolation algorithm. Should be one of:
+            cv2.INTER_NEAREST, cv2.INTER_LINEAR, cv2.INTER_CUBIC, cv2.INTER_AREA, cv2.INTER_LANCZOS4.
+            Default: cv2.INTER_LINEAR.
+        p (float): probability of applying the transform. Default: 1.
+    Targets:
+        image, mask, bboxes
+    Image types:
+        uint8, float32
+    """
+
+    def __init__(self, height, width, erosion_rate=0.0, interpolation=cv2.INTER_LINEAR, always_apply=False, p=1.0):
+        super(RandomSizedBBoxSafeCrop, self).__init__(erosion_rate, always_apply, p)
+        self.height = height
+        self.width = width
+        self.interpolation = interpolation
+
+    def apply(self, img, crop_height=0, crop_width=0, h_start=0, w_start=0, interpolation=cv2.INTER_LINEAR, **params):
+        crop = F.random_crop(img, crop_height, crop_width, h_start, w_start)
+        return FGeometric.resize(crop, self.height, self.width, interpolation)
+
+    def get_transform_init_args_names(self):
+        return super().get_transform_init_args_names() + ("height", "width", "interpolation")
 
 
 class CropAndPad(DualTransform):

diff --git a/tests/test_augmentations.py b/tests/test_augmentations.py
@@ -97,7 +97,7 @@ def test_image_only_augmentations_with_float_values(augmentation_cls, params, fl
             A.CropAndPad: {"px": 10},
             A.Resize: {"height": 10, "width": 10},
         },
-        except_augmentations={A.RandomCropNearBBox, A.RandomSizedBBoxSafeCrop},
+        except_augmentations={A.RandomCropNearBBox, A.RandomSizedBBoxSafeCrop, A.BBoxSafeRandomCrop},
     ),
 )
 def test_dual_augmentations(augmentation_cls, params, image, mask):
@@ -120,7 +120,7 @@ def test_dual_augmentations(augmentation_cls, params, image, mask):
             A.CropAndPad: {"px": 10},
             A.Resize: {"height": 10, "width": 10},
         },
-        except_augmentations={A.RandomCropNearBBox, A.RandomSizedBBoxSafeCrop},
+        except_augmentations={A.RandomCropNearBBox, A.RandomSizedBBoxSafeCrop, A.BBoxSafeRandomCrop},
     ),
 )
 def test_dual_augmentations_with_float_values(augmentation_cls, params, float_image, mask):
@@ -159,7 +159,7 @@ def test_dual_augmentations_with_float_values(augmentation_cls, params, float_im
                 "templates": np.random.randint(low=0, high=256, size=(100, 100, 3), dtype=np.uint8),
             },
         },
-        except_augmentations={A.RandomCropNearBBox, A.RandomSizedBBoxSafeCrop},
+        except_augmentations={A.RandomCropNearBBox, A.RandomSizedBBoxSafeCrop, A.BBoxSafeRandomCrop},
     ),
 )
 def test_augmentations_wont_change_input(augmentation_cls, params, image, mask):
@@ -210,6 +210,7 @@ def test_augmentations_wont_change_input(augmentation_cls, params, image, mask):
             A.RandomToneCurve,
             A.RandomCropNearBBox,
             A.RandomSizedBBoxSafeCrop,
+            A.BBoxSafeRandomCrop,
             A.CropNonEmptyMaskIfExists,
             A.MaskDropout,
         },
@@ -246,6 +247,7 @@ def test_augmentations_wont_change_float_input(augmentation_cls, params, float_i
             A.ISONoise,
             A.RandomCropNearBBox,
             A.RandomSizedBBoxSafeCrop,
+            A.BBoxSafeRandomCrop,
             A.CenterCrop,
             A.Crop,
             A.CropNonEmptyMaskIfExists,
@@ -314,6 +316,7 @@ def test_augmentations_wont_change_shape_grayscale(augmentation_cls, params, ima
         except_augmentations={
             A.RandomCropNearBBox,
             A.RandomSizedBBoxSafeCrop,
+            A.BBoxSafeRandomCrop,
             A.CenterCrop,
             A.Crop,
             A.CropNonEmptyMaskIfExists,
@@ -415,6 +418,7 @@ def test_mask_fill_value(augmentation_cls, params):
             A.RandomRain,
             A.RandomShadow,
             A.RandomSizedBBoxSafeCrop,
+            A.BBoxSafeRandomCrop,
             A.RandomSnow,
             A.RandomSunFlare,
             A.ToFloat,
@@ -473,6 +477,7 @@ def test_multichannel_image_augmentations(augmentation_cls, params):
             A.RandomRain,
             A.RandomShadow,
             A.RandomSizedBBoxSafeCrop,
+            A.BBoxSafeRandomCrop,
             A.RandomSnow,
             A.RandomSunFlare,
             A.ToGray,
@@ -524,6 +529,7 @@ def test_float_multichannel_image_augmentations(augmentation_cls, params):
             A.RandomRain,
             A.RandomShadow,
             A.RandomSizedBBoxSafeCrop,
+            A.BBoxSafeRandomCrop,
             A.RandomSnow,
             A.RandomSunFlare,
             A.ToFloat,
@@ -577,6 +583,7 @@ def test_multichannel_image_augmentations_diff_channels(augmentation_cls, params
             A.RandomRain,
             A.RandomShadow,
             A.RandomSizedBBoxSafeCrop,
+            A.BBoxSafeRandomCrop,
             A.RandomSnow,
             A.RandomSunFlare,
             A.ToGray,
@@ -831,6 +838,7 @@ def test_pixel_domain_adaptation(kind):
             A.CropAndPad: {"px": 10},
             A.Resize: {"height": 10, "width": 10},
             A.RandomSizedBBoxSafeCrop: {"height": 10, "width": 10},
+            A.BBoxSafeRandomCrop: {"erosion_rate": 0.5},
         },
     ),
 )
@@ -846,7 +854,7 @@ def test_non_contiguous_input(augmentation_cls, params, bboxes):
         # requires "cropping_bbox" arg
         aug = augmentation_cls(p=1, **params)
         aug(image=image, mask=mask, cropping_bbox=bboxes[0])
-    elif augmentation_cls == A.RandomSizedBBoxSafeCrop:
+    elif augmentation_cls in [A.RandomSizedBBoxSafeCrop, A.BBoxSafeRandomCrop]:
         # requires "bboxes" arg
         aug = A.Compose([augmentation_cls(p=1, **params)], bbox_params=A.BboxParams(format="pascal_voc"))
         aug(image=image, mask=mask, bboxes=bboxes)

diff --git a/tests/test_serialization.py b/tests/test_serialization.py
@@ -41,6 +41,7 @@
         except_augmentations={
             A.RandomCropNearBBox,
             A.RandomSizedBBoxSafeCrop,
+            A.BBoxSafeRandomCrop,
             A.FDA,
             A.HistogramMatching,
             A.PixelDistributionAdaptation,
@@ -329,6 +330,7 @@ def test_augmentations_serialization(augmentation_cls, params, p, seed, image, m
     A.Lambda,
     A.RandomCropNearBBox,
     A.RandomSizedBBoxSafeCrop,
+    A.BBoxSafeRandomCrop,
     A.GridDropout,
     A.GlassBlur,
     A.TemplateTransform,
@@ -391,6 +393,7 @@ def test_augmentations_serialization_to_file_with_custom_parameters(
             A.CropAndPad: {"px": 10},
             A.Resize: {"height": 10, "width": 10},
             A.RandomSizedBBoxSafeCrop: {"height": 10, "width": 10},
+            A.BBoxSafeRandomCrop: {"erosion_rate": 0.6},
         },
         except_augmentations={
             A.RandomCropNearBBox,
@@ -455,6 +458,7 @@ def test_augmentations_for_bboxes_serialization(
             A.MaskDropout,
             A.OpticalDistortion,
             A.RandomSizedBBoxSafeCrop,
+            A.BBoxSafeRandomCrop,
             A.TemplateTransform,
         },
     ),

diff --git a/tests/test_transforms.py b/tests/test_transforms.py
@@ -156,7 +156,7 @@ def test_elastic_transform_interpolation(monkeypatch, interpolation):
             A.Resize: {"height": 10, "width": 10},
             A.PixelDropout: {"dropout_prob": 0.5, "mask_drop_value": 10, "drop_value": 20},
         },
-        except_augmentations={A.RandomCropNearBBox, A.RandomSizedBBoxSafeCrop, A.PixelDropout},
+        except_augmentations={A.RandomCropNearBBox, A.RandomSizedBBoxSafeCrop, A.BBoxSafeRandomCrop, A.PixelDropout},
     ),
 )
 def test_binary_mask_interpolation(augmentation_cls, params):
@@ -181,7 +181,13 @@ def test_binary_mask_interpolation(augmentation_cls, params):
             A.Resize: {"height": 10, "width": 10},
             A.PixelDropout: {"dropout_prob": 0.5, "mask_drop_value": 10, "drop_value": 20},
         },
-        except_augmentations={A.RandomCropNearBBox, A.RandomSizedBBoxSafeCrop, A.CropAndPad, A.PixelDropout},
+        except_augmentations={
+            A.RandomCropNearBBox,
+            A.RandomSizedBBoxSafeCrop,
+            A.BBoxSafeRandomCrop,
+            A.CropAndPad,
+            A.PixelDropout,
+        },
     ),
 )
 def test_semantic_mask_interpolation(augmentation_cls, params):
@@ -220,6 +226,7 @@ def __test_multiprocessing_support_proc(args):
         except_augmentations={
             A.RandomCropNearBBox,
             A.RandomSizedBBoxSafeCrop,
+            A.BBoxSafeRandomCrop,
             A.CropNonEmptyMaskIfExists,
             A.FDA,
             A.HistogramMatching,