use nearest interpolation when applying segmentation warp for affine augmentation

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

Switch to using nearest pixel interpolation when warping and added unit test.

Reviewed By: wat3rBro

Differential Revision: D41042506

fbshipit-source-id: 92b817f21e862422428a0d0df969ec9e037f99fb

d2go/data/transforms/affine.py

 #!/usr/bin/env python3
 # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
 
-
+import copy
 import json
 import random
 from typing import List, Optional, Tuple
@@ -37,18 +37,29 @@ class AffineTransform(Transform):
         if border_mode is not None:
             self.warp_kwargs["borderMode"] = border_mode
 
-    def apply_image(self, img: np.ndarray) -> np.ndarray:
+    def _warp_array(self, input_data: np.array, interp_flag: Optional[int] = None):
+        warp_kwargs = copy.deepcopy(self.warp_kwargs)
+
+        if interp_flag is not None:
+            flags = warp_kwargs.get("flags", 0)
+            # remove previous interp and add the new one
+            flags = (flags - (flags & cv2.INTER_MAX)) + interp_flag
+            warp_kwargs["flags"] = flags
+
         M = self.M
         if self.is_inversed_M:
             M = M[:2]
         img = cv2.warpAffine(
-            img,
+            input_data,
             M,
             (int(self.img_w), (self.img_h)),
-            **self.warp_kwargs,
+            **warp_kwargs,
         )
         return img
 
+    def apply_image(self, img: np.ndarray) -> np.ndarray:
+        return self._warp_array(img)
+
     def apply_coords(self, coords: np.ndarray) -> np.ndarray:
         # Add row of ones to enable matrix multiplication
         coords = coords.T
@@ -61,7 +72,7 @@ class AffineTransform(Transform):
         return coords
 
     def apply_segmentation(self, img: np.ndarray) -> np.ndarray:
-        raise NotImplementedError()
+        return self._warp_array(img, interp_flag=cv2.INTER_NEAREST)
 
 
 class RandomPivotScaling(TransformGen):

tests/data/test_data_transforms_affine.py

@@ -86,9 +86,25 @@ def generate_test_data(
         borderMode=cv2.BORDER_REPLICATE,
     )
 
-    # Create test boxes
-    M_inv = np.vstack([M_inv, [0.0, 0.0, 1.0]])
+    # Create annotations
     test_bbox = [0.25 * img_w, 0.25 * img_h, 0.75 * img_h, 0.75 * img_h]
+
+    # Generate segmentation test data
+    segm_mask = np.zeros_like(source_img)
+    segm_mask[
+        int(test_bbox[0]) : int(test_bbox[2]), int(test_bbox[1]) : int(test_bbox[3])
+    ] = 255
+
+    exp_out_segm = cv2.warpAffine(
+        segm_mask,
+        M_inv,
+        (out_w, out_h),
+        flags=cv2.WARP_INVERSE_MAP + cv2.INTER_NEAREST,
+        borderMode=cv2.BORDER_REPLICATE,
+    )
+
+    # Generate bounding box test data
+    M_inv = np.vstack([M_inv, [0.0, 0.0, 1.0]])
     points = np.array(
         [
             [test_bbox[0], test_bbox[0], test_bbox[2], test_bbox[2]],
@@ -97,7 +113,12 @@ def generate_test_data(
     ).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])
+
+    return (
+        aug_str,
+        AugInput(source_img, boxes=[test_bbox], sem_seg=segm_mask),
+        (exp_out_img, [out_bbox], exp_out_segm),
+    )
 
 
 def warp_points(coords: np.array, xfm_M: np.array):
@@ -110,17 +131,25 @@ def warp_points(coords: np.array, xfm_M: np.array):
 
 
 class TestDataTransformsAffine(unittest.TestCase):
-    def _check_array_close(self, aug_output, exp_img, exp_bboxes):
+    def _validate_results(self, aug_output, exp_outputs):
+        exp_img = exp_outputs[0]
         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]} ",
         )
 
+        exp_bboxes = exp_outputs[1]
         self.assertTrue(
             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]} ",
         )
 
+        exp_segm = exp_outputs[2]
+        self.assertTrue(
+            np.allclose(exp_segm, aug_output.sem_seg),
+            f"Augmented segm not the same, expecting\n{exp_segm} \n   got\n{aug_output.sem_seg[:,:]} ",
+        )
+
     def test_affine_transforms_angle(self):
         default_cfg = Detectron2GoRunner.get_default_cfg()
 
@@ -129,15 +158,13 @@ class TestDataTransformsAffine(unittest.TestCase):
         img[((img_sz + 1) // 2) - 1, :, :] = 255
 
         for angle in [45, 90]:
-            aug_str, aug_input, (exp_out_img, exp_out_bboxes) = generate_test_data(
-                img, angle=angle
-            )
+            aug_str, aug_input, exp_outputs = generate_test_data(img, angle=angle)
             default_cfg.D2GO_DATA.AUG_OPS.TRAIN = [aug_str]
             tfm = build_transform_gen(default_cfg, is_train=True)
 
             # Test augmentation
             aug_output, _ = apply_augmentations(tfm, aug_input)
-            self._check_array_close(aug_output, exp_out_img, exp_out_bboxes)
+            self._validate_results(aug_output, exp_outputs)
 
     def test_affine_transforms_translation(self):
         default_cfg = Detectron2GoRunner.get_default_cfg()
@@ -148,7 +175,7 @@ class TestDataTransformsAffine(unittest.TestCase):
 
         for translation in [0, 1, 2]:
             # Test image
-            aug_str, aug_input, (exp_out_img, exp_out_bboxes) = generate_test_data(
+            aug_str, aug_input, exp_outputs = generate_test_data(
                 img, translation=translation
             )
             default_cfg.D2GO_DATA.AUG_OPS.TRAIN = [aug_str]
@@ -156,7 +183,7 @@ class TestDataTransformsAffine(unittest.TestCase):
 
             # Test augmentation
             aug_output, _ = apply_augmentations(tfm, aug_input)
-            self._check_array_close(aug_output, exp_out_img, exp_out_bboxes)
+            self._validate_results(aug_output, exp_outputs)
 
     def test_affine_transforms_shear(self):
         default_cfg = Detectron2GoRunner.get_default_cfg()
@@ -166,15 +193,13 @@ class TestDataTransformsAffine(unittest.TestCase):
         img[((img_sz + 1) // 2) - 1, :, :] = 255
 
         for shear in [0, 1, 2]:
-            aug_str, aug_input, (exp_out_img, exp_out_bboxes) = generate_test_data(
-                img, shear=shear
-            )
+            aug_str, aug_input, exp_outputs = generate_test_data(img, shear=shear)
             default_cfg.D2GO_DATA.AUG_OPS.TRAIN = [aug_str]
             tfm = build_transform_gen(default_cfg, is_train=True)
 
             # Test augmentation
             aug_output, _ = apply_augmentations(tfm, aug_input)
-            self._check_array_close(aug_output, exp_out_img, exp_out_bboxes)
+            self._validate_results(aug_output, exp_outputs)
 
     def test_affine_transforms_scale(self):
         default_cfg = Detectron2GoRunner.get_default_cfg()
@@ -184,15 +209,13 @@ class TestDataTransformsAffine(unittest.TestCase):
         img[((img_sz + 1) // 2) - 1, :, :] = 255
 
         for scale in [0.9, 1, 1.1]:
-            aug_str, aug_input, (exp_out_img, exp_out_bboxes) = generate_test_data(
-                img, scale=scale
-            )
+            aug_str, aug_input, exp_outputs = generate_test_data(img, scale=scale)
             default_cfg.D2GO_DATA.AUG_OPS.TRAIN = [aug_str]
             tfm = build_transform_gen(default_cfg, is_train=True)
 
             # Test augmentation
             aug_output, _ = apply_augmentations(tfm, aug_input)
-            self._check_array_close(aug_output, exp_out_img, exp_out_bboxes)
+            self._validate_results(aug_output, exp_outputs)
 
     def test_affine_transforms_angle_non_square(self):
         default_cfg = Detectron2GoRunner.get_default_cfg()
@@ -202,7 +225,7 @@ class TestDataTransformsAffine(unittest.TestCase):
         img[((img_sz + 1) // 2) - 1, :, :] = 255
 
         for keep_aspect_ratio in [False, True]:
-            aug_str, aug_input, (exp_out_img, exp_out_bboxes) = generate_test_data(
+            aug_str, aug_input, exp_outputs = generate_test_data(
                 img, angle=45, keep_aspect_ratio=keep_aspect_ratio
             )
             default_cfg.D2GO_DATA.AUG_OPS.TRAIN = [aug_str]
@@ -210,7 +233,7 @@ class TestDataTransformsAffine(unittest.TestCase):
 
             # Test augmentation
             aug_output, _ = apply_augmentations(tfm, aug_input)
-            self._check_array_close(aug_output, exp_out_img, exp_out_bboxes)
+            self._validate_results(aug_output, exp_outputs)
 
     def test_affine_transforms_angle_no_fit_to_frame(self):
         default_cfg = Detectron2GoRunner.get_default_cfg()
@@ -219,7 +242,7 @@ class TestDataTransformsAffine(unittest.TestCase):
         img = np.zeros((img_sz, img_sz, 3)).astype(np.uint8)
         img[((img_sz + 1) // 2) - 1, :, :] = 255
 
-        aug_str, aug_input, (exp_out_img, exp_out_bboxes) = generate_test_data(
+        aug_str, aug_input, exp_outputs = generate_test_data(
             img, angle=45, fit_in_frame=False
         )
         default_cfg.D2GO_DATA.AUG_OPS.TRAIN = [aug_str]
@@ -227,4 +250,4 @@ class TestDataTransformsAffine(unittest.TestCase):
 
         # Test augmentation
         aug_output, _ = apply_augmentations(tfm, aug_input)
-        self._check_array_close(aug_output, exp_out_img, exp_out_bboxes)
+        self._validate_results(aug_output, exp_outputs)