Skip to content

GitLab

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

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
parent 42ac339f
Show whitespace changes
Inline Side-by-side
Showing with 69 additions and 30 deletions
d2go/data/transforms/affine.py
View file @ 69dac5b5
......@@ -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
View file @ 69dac5b5
......@@ -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_output, test_output),
f"Augmented image not the same, expecting\n{exp_output[:,:,0]} \n got\n{test_output[:,:,0]} ",
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_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)
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