fix affine augmentation coordinate warping

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

1. Fix bounding box coordinate warping
2. Add apply_segmentation warning (will followup in another fix)

Reviewed By: wat3rBro

Differential Revision: D41013775

fbshipit-source-id: 3652b04c1622fe35fa9893dc22350f7d59b37c6e

d2go/data/transforms/affine.py

@@ -60,6 +60,9 @@ class AffineTransform(Transform):
         coords = (M @ coords)[:2, :].T
         return coords
 
+    def apply_segmentation(self, img: np.ndarray) -> np.ndarray:
+        raise NotImplementedError()
+
 
 class RandomPivotScaling(TransformGen):
     """
@@ -211,6 +214,7 @@ class RandomAffine(TransformGen):
             center, angle, translation, scale, shear
         )
         M_inv = np.array(M_inv).reshape((2, 3))
+        M_inv = np.vstack([M_inv, [0.0, 0.0, 1.0]])
 
         return AffineTransform(
             M_inv,

tests/data/test_data_transforms_affine.py

@@ -10,7 +10,7 @@ import numpy as np
 import torchvision.transforms as T
 from d2go.data.transforms.build import build_transform_gen
 from d2go.runner import Detectron2GoRunner
-from detectron2.data.transforms import apply_augmentations
+from detectron2.data.transforms import apply_augmentations, AugInput
 
 
 def generate_test_data(
@@ -78,7 +78,6 @@ def generate_test_data(
         center, angle, translation_t, scale, [shear, shear]
     )
     M_inv = np.array(M_inv).reshape((2, 3))
-
     exp_out_img = cv2.warpAffine(
         source_img,
         M_inv,
@@ -87,14 +86,39 @@ def generate_test_data(
         borderMode=cv2.BORDER_REPLICATE,
     )
 
-    return aug_str, exp_out_img
+    # Create test boxes
+    M_inv = np.vstack([M_inv, [0.0, 0.0, 1.0]])
+    test_bbox = [0.25 * img_w, 0.25 * img_h, 0.75 * img_h, 0.75 * img_h]
+    points = np.array(
+        [
+            [test_bbox[0], test_bbox[0], test_bbox[2], test_bbox[2]],
+            [test_bbox[1], test_bbox[3], test_bbox[1], test_bbox[3]],
+        ]
+    ).T
+    _xp = warp_points(points, M_inv)
+    out_bbox = [min(_xp[:, 0]), min(_xp[:, 1]), max(_xp[:, 0]), max(_xp[:, 1])]
+    return aug_str, AugInput(source_img, boxes=[test_bbox]), (exp_out_img, [out_bbox])
+
+
+def warp_points(coords: np.array, xfm_M: np.array):
+    coords = coords.T
+    ones = np.ones((1, coords.shape[1]))
+    coords = np.vstack((coords, ones))
+    M = np.linalg.inv(xfm_M)
+    coords = (M @ coords)[:2, :].T
+    return coords
 
 
 class TestDataTransformsAffine(unittest.TestCase):
-    def _check_array_close(self, test_output, exp_output):
+    def _check_array_close(self, aug_output, exp_img, exp_bboxes):
+        self.assertTrue(
+            np.allclose(exp_img, aug_output.image),
+            f"Augmented image not the same, expecting\n{exp_img[:,:,0]} \n   got\n{aug_output.image[:,:,0]} ",
+        )
+
         self.assertTrue(
-            np.allclose(exp_output, test_output),
-            f"Augmented image not the same, expecting\n{exp_output[:,:,0]} \n   got\n{test_output[:,:,0]} ",
+            np.allclose(exp_bboxes, aug_output.boxes, atol=0.000001),
+            f"Augmented bbox not the same, expecting\n{exp_img[:,:,0]} \n   got\n{aug_output.image[:,:,0]} ",
         )
 
     def test_affine_transforms_angle(self):
@@ -105,13 +129,15 @@ class TestDataTransformsAffine(unittest.TestCase):
         img[((img_sz + 1) // 2) - 1, :, :] = 255
 
         for angle in [45, 90]:
-            aug_str, exp_out_img = generate_test_data(img, angle=angle)
-
+            aug_str, aug_input, (exp_out_img, exp_out_bboxes) = generate_test_data(
+                img, angle=angle
+            )
             default_cfg.D2GO_DATA.AUG_OPS.TRAIN = [aug_str]
             tfm = build_transform_gen(default_cfg, is_train=True)
-            trans_img, _ = apply_augmentations(tfm, img)
 
-            self._check_array_close(trans_img, exp_out_img)
+            # Test augmentation
+            aug_output, _ = apply_augmentations(tfm, aug_input)
+            self._check_array_close(aug_output, exp_out_img, exp_out_bboxes)
 
     def test_affine_transforms_translation(self):
         default_cfg = Detectron2GoRunner.get_default_cfg()
@@ -121,13 +147,16 @@ class TestDataTransformsAffine(unittest.TestCase):
         img[((img_sz + 1) // 2) - 1, :, :] = 255
 
         for translation in [0, 1, 2]:
-            aug_str, exp_out_img = generate_test_data(img, translation=translation)
-
+            # Test image
+            aug_str, aug_input, (exp_out_img, exp_out_bboxes) = generate_test_data(
+                img, translation=translation
+            )
             default_cfg.D2GO_DATA.AUG_OPS.TRAIN = [aug_str]
             tfm = build_transform_gen(default_cfg, is_train=True)
-            trans_img, _ = apply_augmentations(tfm, img)
 
-            self._check_array_close(trans_img, exp_out_img)
+            # Test augmentation
+            aug_output, _ = apply_augmentations(tfm, aug_input)
+            self._check_array_close(aug_output, exp_out_img, exp_out_bboxes)
 
     def test_affine_transforms_shear(self):
         default_cfg = Detectron2GoRunner.get_default_cfg()
@@ -137,13 +166,15 @@ class TestDataTransformsAffine(unittest.TestCase):
         img[((img_sz + 1) // 2) - 1, :, :] = 255
 
         for shear in [0, 1, 2]:
-            aug_str, exp_out_img = generate_test_data(img, shear=shear)
-
+            aug_str, aug_input, (exp_out_img, exp_out_bboxes) = generate_test_data(
+                img, shear=shear
+            )
             default_cfg.D2GO_DATA.AUG_OPS.TRAIN = [aug_str]
             tfm = build_transform_gen(default_cfg, is_train=True)
-            trans_img, _ = apply_augmentations(tfm, img)
 
-            self._check_array_close(trans_img, exp_out_img)
+            # Test augmentation
+            aug_output, _ = apply_augmentations(tfm, aug_input)
+            self._check_array_close(aug_output, exp_out_img, exp_out_bboxes)
 
     def test_affine_transforms_scale(self):
         default_cfg = Detectron2GoRunner.get_default_cfg()
@@ -153,13 +184,15 @@ class TestDataTransformsAffine(unittest.TestCase):
         img[((img_sz + 1) // 2) - 1, :, :] = 255
 
         for scale in [0.9, 1, 1.1]:
-            aug_str, exp_out_img = generate_test_data(img, scale=scale)
-
+            aug_str, aug_input, (exp_out_img, exp_out_bboxes) = generate_test_data(
+                img, scale=scale
+            )
             default_cfg.D2GO_DATA.AUG_OPS.TRAIN = [aug_str]
             tfm = build_transform_gen(default_cfg, is_train=True)
-            trans_img, _ = apply_augmentations(tfm, img)
 
-            self._check_array_close(trans_img, exp_out_img)
+            # Test augmentation
+            aug_output, _ = apply_augmentations(tfm, aug_input)
+            self._check_array_close(aug_output, exp_out_img, exp_out_bboxes)
 
     def test_affine_transforms_angle_non_square(self):
         default_cfg = Detectron2GoRunner.get_default_cfg()
@@ -169,15 +202,15 @@ class TestDataTransformsAffine(unittest.TestCase):
         img[((img_sz + 1) // 2) - 1, :, :] = 255
 
         for keep_aspect_ratio in [False, True]:
-            aug_str, exp_out_img = generate_test_data(
+            aug_str, aug_input, (exp_out_img, exp_out_bboxes) = generate_test_data(
                 img, angle=45, keep_aspect_ratio=keep_aspect_ratio
             )
-
             default_cfg.D2GO_DATA.AUG_OPS.TRAIN = [aug_str]
             tfm = build_transform_gen(default_cfg, is_train=True)
-            trans_img, _ = apply_augmentations(tfm, img)
 
-            self._check_array_close(trans_img, exp_out_img)
+            # Test augmentation
+            aug_output, _ = apply_augmentations(tfm, aug_input)
+            self._check_array_close(aug_output, exp_out_img, exp_out_bboxes)
 
     def test_affine_transforms_angle_no_fit_to_frame(self):
         default_cfg = Detectron2GoRunner.get_default_cfg()
@@ -186,10 +219,12 @@ class TestDataTransformsAffine(unittest.TestCase):
         img = np.zeros((img_sz, img_sz, 3)).astype(np.uint8)
         img[((img_sz + 1) // 2) - 1, :, :] = 255
 
-        aug_str, exp_out_img = generate_test_data(img, angle=45, fit_in_frame=False)
-
+        aug_str, aug_input, (exp_out_img, exp_out_bboxes) = generate_test_data(
+            img, angle=45, fit_in_frame=False
+        )
         default_cfg.D2GO_DATA.AUG_OPS.TRAIN = [aug_str]
         tfm = build_transform_gen(default_cfg, is_train=True)
-        trans_img, _ = apply_augmentations(tfm, img)
 
-        self._check_array_close(trans_img, exp_out_img)
+        # Test augmentation
+        aug_output, _ = apply_augmentations(tfm, aug_input)
+        self._check_array_close(aug_output, exp_out_img, exp_out_bboxes)