Merge pull request #139 from LinaShiryaeva/master

add RandomCenteredCrop augmentation that depends on the input parameters

albumentations/augmentations/functional.py

@@ -221,6 +221,19 @@ def random_crop(img, crop_height, crop_width, h_start, w_start):
     return img
 
 
+def clamping_crop(img, x_min, y_min, x_max, y_max):
+    h, w = img.shape[:2]
+    if x_min < 0:
+        x_min = 0
+    if y_min < 0:
+        y_min = 0
+    if y_max >= h:
+        y_max = h - 1
+    if x_max >= w:
+        x_max = w - 1
+    return img[int(y_min):int(y_max), int(x_min):int(x_max)]
+
+
 def shift_hsv(img, hue_shift, sat_shift, val_shift):
     dtype = img.dtype
     img = cv2.cvtColor(img, cv2.COLOR_RGB2HSV)

albumentations/augmentations/transforms.py

@@ -14,7 +14,7 @@
            'ElasticTransform', 'HueSaturationValue', 'PadIfNeeded', 'RGBShift', 'RandomBrightness', 'RandomContrast',
            'MotionBlur', 'MedianBlur', 'GaussNoise', 'CLAHE', 'ChannelShuffle', 'InvertImg', 'ToGray',
            'JpegCompression', 'Cutout', 'ToFloat', 'FromFloat', 'Crop', 'RandomScale', 'LongestMaxSize',
-           'SmallestMaxSize', 'Resize', 'RandomSizedCrop', 'RandomBrightnessContrast']
+           'SmallestMaxSize', 'Resize', 'RandomSizedCrop', 'RandomBrightnessContrast', 'RandomCropNearBBox']
 
 
 class PadIfNeeded(DualTransform):
@@ -460,6 +460,54 @@ def apply_to_bbox(self, bbox, **params):
         return F.bbox_random_crop(bbox, self.height, self.width, **params)
 
 
+class RandomCropNearBBox(DualTransform):
+    """Crop bbox from image with random shift by x,y coordinates
+
+    Args:
+        max_part_shift (float): float value in (0.0, 1.0) range. Default 0.3
+        p (float): probability of applying the transform. Default: 1.
+
+    Targets:
+        image
+
+    Image types:
+        uint8, float32
+    """
+
+    def __init__(self, max_part_shift=0.3, always_apply=False, p=1.0):
+        super(RandomCropNearBBox, self).__init__(always_apply, p)
+        self.max_part_shift = max_part_shift
+
+    def apply(self, img, x_min=0, x_max=0, y_min=0, y_max=0, **params):
+        return F.clamping_crop(img, x_min, x_max, y_min, y_max)
+
+    def get_params_dependent_on_targets(self, params):
+        bbox = params['cropping_bbox']
+        h_max_shift = int((bbox[3] - bbox[1]) * self.max_part_shift)
+        w_max_shift = int((bbox[2] - bbox[0]) * self.max_part_shift)
+
+        x_min = bbox[0] - random.randint(-w_max_shift, w_max_shift)
+        x_max = bbox[2] + random.randint(-w_max_shift, w_max_shift)
+
+        y_min = bbox[1] - random.randint(-h_max_shift, h_max_shift)
+        y_max = bbox[3] + random.randint(-h_max_shift, h_max_shift)
+
+        return {'x_min': x_min,
+                'x_max': x_max,
+                'y_min': y_min,
+                'y_max': y_max
+                }
+
+    def apply_to_bbox(self, bbox, x_min=0, x_max=0, y_min=0, y_max=0, **params):
+        h_start = y_min
+        w_start = x_min
+        return F.bbox_crop(bbox, y_max - y_min, x_max - x_min, h_start, w_start, **params)
+
+    @property
+    def targets_as_params(self):
+        return ['cropping_bbox']
+
+
 class RandomSizedCrop(DualTransform):
     """Crop a random part of the input and rescale it to some size.
 

tests/test_augmentations.py

@@ -14,7 +14,7 @@
     Rotate, ShiftScaleRotate, CenterCrop, OpticalDistortion, GridDistortion, ElasticTransform, ToGray, RandomGamma, \
     JpegCompression, HueSaturationValue, RGBShift, RandomBrightness, RandomContrast, Blur, MotionBlur, MedianBlur, \
     GaussNoise, CLAHE, ChannelShuffle, InvertImg, IAAEmboss, IAASuperpixels, IAASharpen, IAAAdditiveGaussianNoise, \
-    IAAPiecewiseAffine, IAAPerspective, Cutout, Normalize, ToFloat, FromFloat, RandomSizedCrop
+    IAAPiecewiseAffine, IAAPerspective, Cutout, Normalize, ToFloat, FromFloat, RandomSizedCrop, RandomCropNearBBox
 
 
 @pytest.mark.parametrize(['augmentation_cls', 'params'], [
@@ -317,3 +317,13 @@ def test_augmentations_wont_change_shape_rgb(augmentation_cls, params, image, ma
     result = aug(image=image_3ch, mask=mask_3ch)
     assert np.array_equal(image_3ch.shape, result['image'].shape)
     assert np.array_equal(mask_3ch.shape, result['mask'].shape)
+
+
+@pytest.mark.parametrize(['augmentation_cls', 'params'], [
+    [RandomCropNearBBox, {'max_part_shift': 0.15}],
+])
+def test_image_only_crop_around_bbox_augmentation(augmentation_cls, params, image, mask):
+    aug = augmentation_cls(p=1, **params)
+    annotations = {'image': image, 'cropping_bbox': [-59, 77, 177, 231]}
+    data = aug(**annotations)
+    assert data['image'].dtype == np.uint8