Skip to content

GitLab

Commit d660903c authored by Sam Tsai's avatar Sam Tsai Committed by Facebook GitHub Bot
Browse files

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
parent 69dac5b5
Show whitespace changes
Inline Side-by-side
Showing with 61 additions and 27 deletions
d2go/data/transforms/affine.py
View file @ d660903c
#!/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
View file @ d660903c
......@@ -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)
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment