test_data_transforms_affine.py

#!/usr/bin/env python3
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved

import json
import unittest
from typing import Tuple

import cv2
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, AugInput


def generate_test_data(
    source_img: np.ndarray,
    angle: float = 0,
    translation: float = 0,
    scale: float = 1,
    shear: float = 0,
    fit_in_frame: bool = True,
    keep_aspect_ratio: bool = False,
) -> Tuple[str, np.ndarray]:
    # Augmentation dictionary
    aug_dict = {
        "prob": 1.0,
        "angle_range": [angle, angle],
        "translation_range": [translation, translation],
        "scale_range": [scale, scale],
        "shear_range": [shear, shear],
        "keep_aspect_ratio": keep_aspect_ratio,
        "fit_in_frame": fit_in_frame,
    }
    aug_str = "RandomAffineOp::" + json.dumps(aug_dict)

    # Get image info
    img_h, img_w = source_img.shape[0:2]
    center = [img_w / 2, img_h / 2]

    # Compute output_size
    max_size = max(img_w, img_h)
    out_w, out_h = (img_w, img_h) if keep_aspect_ratio else (max_size, max_size)

    if fit_in_frame:
        # Warp once to figure scale adjustment
        M_inv = T.functional._get_inverse_affine_matrix(
            center, angle, [0, 0], 1, [shear, shear]
        )
        M_inv.extend([0.0, 0.0, 1.0])
        M_inv = np.array(M_inv).reshape((3, 3))
        M = np.linalg.inv(M_inv)

        # Center in output patch
        img_corners = np.array(
            [
                [0, 0, img_w - 1, img_w - 1],
                [0, img_h - 1, 0, img_h - 1],
                [1, 1, 1, 1],
            ]
        )
        new_corners = M @ img_corners
        x_range = np.ceil(np.amax(new_corners[0]) - np.amin(new_corners[0]))
        y_range = np.ceil(np.amax(new_corners[1]) - np.amin(new_corners[1]))

        # Apply translation and scale after centering in output patch
        scale_adjustment = min(out_w / x_range, out_h / y_range)
        scale *= scale_adjustment

    # Adjust output center location
    translation_t = [translation, translation]
    translation_adjustment = [(out_w - img_w) / 2, (out_h - img_h) / 2]
    translation_t[0] += translation_adjustment[0]
    translation_t[1] += translation_adjustment[1]

    # Test data output generation
    M_inv = T.functional._get_inverse_affine_matrix(
        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,
        (out_w, out_h),
        flags=cv2.WARP_INVERSE_MAP + cv2.INTER_LINEAR,
        borderMode=cv2.BORDER_REPLICATE,
    )

    # 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]],
            [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], sem_seg=segm_mask),
        (exp_out_img, [out_bbox], exp_out_segm),
    )


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 _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()

        img_sz = 11
        img = np.zeros((img_sz, img_sz, 3)).astype(np.uint8)
        img[((img_sz + 1) // 2) - 1, :, :] = 255

        for angle in [45, 90]:
            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._validate_results(aug_output, exp_outputs)

    def test_affine_transforms_translation(self):
        default_cfg = Detectron2GoRunner.get_default_cfg()

        img_sz = 11
        img = np.zeros((img_sz, img_sz, 3)).astype(np.uint8)
        img[((img_sz + 1) // 2) - 1, :, :] = 255

        for translation in [0, 1, 2]:
            # Test image
            aug_str, aug_input, exp_outputs = generate_test_data(
                img, translation=translation
            )
            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._validate_results(aug_output, exp_outputs)

    def test_affine_transforms_shear(self):
        default_cfg = Detectron2GoRunner.get_default_cfg()

        img_sz = 11
        img = np.zeros((img_sz, img_sz, 3)).astype(np.uint8)
        img[((img_sz + 1) // 2) - 1, :, :] = 255

        for shear in [0, 1, 2]:
            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._validate_results(aug_output, exp_outputs)

    def test_affine_transforms_scale(self):
        default_cfg = Detectron2GoRunner.get_default_cfg()

        img_sz = 11
        img = np.zeros((img_sz, img_sz, 3)).astype(np.uint8)
        img[((img_sz + 1) // 2) - 1, :, :] = 255

        for scale in [0.9, 1, 1.1]:
            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._validate_results(aug_output, exp_outputs)

    def test_affine_transforms_angle_non_square(self):
        default_cfg = Detectron2GoRunner.get_default_cfg()

        img_sz = 11
        img = np.zeros((img_sz, img_sz - 2, 3)).astype(np.uint8)
        img[((img_sz + 1) // 2) - 1, :, :] = 255

        for keep_aspect_ratio in [False, True]:
            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]
            tfm = build_transform_gen(default_cfg, is_train=True)

            # Test augmentation
            aug_output, _ = apply_augmentations(tfm, aug_input)
            self._validate_results(aug_output, exp_outputs)

    def test_affine_transforms_angle_no_fit_to_frame(self):
        default_cfg = Detectron2GoRunner.get_default_cfg()

        img_sz = 11
        img = np.zeros((img_sz, img_sz, 3)).astype(np.uint8)
        img[((img_sz + 1) // 2) - 1, :, :] = 255

        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]
        tfm = build_transform_gen(default_cfg, is_train=True)

        # Test augmentation
        aug_output, _ = apply_augmentations(tfm, aug_input)
        self._validate_results(aug_output, exp_outputs)