add an augmentation to pad image to square.

Summary:
Pull Request resolved: https://github.com/facebookresearch/d2go/pull/428

add an augmentation to pad image to square.
* For example, image with shape (10, 7, 3) will become (10, 10, 3) and pad with value specified by `pad_value`.

Reviewed By: tax313, wat3rBro

Differential Revision: D41545182

fbshipit-source-id: 6d5fd9d16984a9904d44f22386920cdf130edda7

d2go/data/transforms/crop.py

@@ -42,6 +42,7 @@ class PadTransform(Transform):
         org_w: int,
         org_h: int,
         pad_mode: str = "constant",
+        pad_value: float = 0.0,
     ):
         super().__init__()
         assert x0 + w <= org_w
@@ -59,7 +60,7 @@ class PadTransform(Transform):
         ]
         pad_args = {"mode": self.pad_mode}
         if self.pad_mode == "constant":
-            pad_args["constant_values"] = 0
+            pad_args["constant_values"] = self.pad_value
         ret = np.pad(img, pad_width=tuple(pad_width), **pad_args)
         return ret
 
@@ -88,6 +89,35 @@ class PadBorderDivisible(aug.Augmentation):
         return PadTransform(0, 0, W, H, new_w, new_h, pad_mode=self.pad_mode)
 
 
+class PadToSquare(aug.Augmentation):
+    """Pad the image to square"""
+
+    def __init__(
+        self,
+        pad_mode: str = "constant",
+        pad_value: float = 0.0,
+    ):
+        super().__init__()
+        self.pad_mode = pad_mode
+        self.pad_value = pad_value
+
+    def get_transform(self, image: np.ndarray) -> Transform:
+        """image: HxWxC"""
+        assert len(image.shape) == 3 and image.shape[2] in [1, 3]
+        H, W = image.shape[:2]
+        new_h = new_w = max(H, W)
+        return PadTransform(
+            0,
+            0,
+            W,
+            H,
+            new_w,
+            new_h,
+            pad_mode=self.pad_mode,
+            pad_value=self.pad_value,
+        )
+
+
 class RandomCropFixedAspectRatio(aug.Augmentation):
     def __init__(
         self,
@@ -170,6 +200,16 @@ def CropBoundaryOp(
     return [CropBoundary(**kwargs)]
 
 
+# example repr: 'PadToSquareOp::{"pad_value": 255.0}'
+@TRANSFORM_OP_REGISTRY.register()
+def PadToSquareOp(
+    cfg: CfgNode, arg_str: str, is_train: bool
+) -> List[Union[aug.Augmentation, Transform]]:
+    kwargs = _json_load(arg_str) if arg_str is not None else {}
+    assert isinstance(kwargs, dict)
+    return [PadToSquare(**kwargs)]
+
+
 # example repr: "RandomCropFixedAspectRatioOp::{'crop_aspect_ratios_list': [0.5], 'scale_range': [0.8, 1.2], 'offset_scale_range': [-0.3, 0.3]}"
 @TRANSFORM_OP_REGISTRY.register()
 def RandomCropFixedAspectRatioOp(

tests/data/test_data_transforms_crop.py

@@ -139,7 +139,7 @@ class TestDataTransformsCrop(unittest.TestCase):
             trans.inverse().inverse().apply_image(full_img_gt), crop_img_gt
         )
 
-    def test_pad_transform(self):
+    def test_pad_border_divisible_transform(self):
         img_h, img_w = 10, 7
         divisibility = 8
 
@@ -160,6 +160,16 @@ class TestDataTransformsCrop(unittest.TestCase):
         self.assertEqual(inverse_mask.shape, (10, 7))
         self.assertArrayEqual(mask, inverse_mask)
 
+    def test_pad_to_square_augmentation(self):
+        img_h, img_w = 5, 3
+
+        aug = tf_crop.PadToSquare(pad_value=255)
+
+        img = np.ones([img_h, img_w, 3])
+        trans = aug.get_transform(img)
+        pad_img = trans.apply_image(img)
+        self.assertEqual(pad_img.shape, (5, 5, 3))
+
     def test_random_instance_crop(self):
         from detectron2.data import detection_utils as du
         from detectron2.data.transforms.augmentation import AugInput, AugmentationList