detectron2/tests/data/test_transforms.py

# -*- coding: utf-8 -*-
# Copyright (c) Facebook, Inc. and its affiliates.

import logging
import numpy as np
import unittest
from unittest import mock
import torch
from PIL import Image, ImageOps
from torch.nn import functional as F

from detectron2.config import get_cfg
from detectron2.data import detection_utils
from detectron2.data import transforms as T
from detectron2.utils.logger import setup_logger

logger = logging.getLogger(__name__)


def polygon_allclose(poly1, poly2):
    """
    Test whether two polygons are the same.
    Both arguments are nx2 numpy arrays.
    """
    # ABCD and CDAB are the same polygon. So it's important to check after rolling
    for k in range(len(poly1)):
        rolled_poly1 = np.roll(poly1, k, axis=0)
        if np.allclose(rolled_poly1, poly2):
            return True
    return False


class TestTransforms(unittest.TestCase):
    def setUp(self):
        setup_logger()

    def test_apply_rotated_boxes(self):
        np.random.seed(125)
        cfg = get_cfg()
        is_train = True
        augs = detection_utils.build_augmentation(cfg, is_train)
        image = np.random.rand(200, 300)
        image, transforms = T.apply_augmentations(augs, image)
        image_shape = image.shape[:2]  # h, w
        assert image_shape == (800, 1200)
        annotation = {"bbox": [179, 97, 62, 40, -56]}

        boxes = np.array([annotation["bbox"]], dtype=np.float64)  # boxes.shape = (1, 5)
        transformed_bbox = transforms.apply_rotated_box(boxes)[0]

        expected_bbox = np.array([484, 388, 248, 160, 56], dtype=np.float64)
        err_msg = "transformed_bbox = {}, expected {}".format(transformed_bbox, expected_bbox)
        assert np.allclose(transformed_bbox, expected_bbox), err_msg

    def test_resize_and_crop(self):
        np.random.seed(125)
        min_scale = 0.2
        max_scale = 2.0
        target_height = 1100
        target_width = 1000
        resize_aug = T.ResizeScale(min_scale, max_scale, target_height, target_width)
        fixed_size_crop_aug = T.FixedSizeCrop((target_height, target_width))
        hflip_aug = T.RandomFlip()
        augs = [resize_aug, fixed_size_crop_aug, hflip_aug]
        original_image = np.random.rand(900, 800)
        image, transforms = T.apply_augmentations(augs, original_image)
        image_shape = image.shape[:2]  # h, w
        self.assertEqual((1100, 1000), image_shape)

        boxes = np.array(
            [[91, 46, 144, 111], [523, 251, 614, 295]],
            dtype=np.float64,
        )
        transformed_bboxs = transforms.apply_box(boxes)
        expected_bboxs = np.array(
            [
                [895.42, 33.42666667, 933.91125, 80.66],
                [554.0825, 182.39333333, 620.17125, 214.36666667],
            ],
            dtype=np.float64,
        )
        err_msg = "transformed_bbox = {}, expected {}".format(transformed_bboxs, expected_bboxs)
        self.assertTrue(np.allclose(transformed_bboxs, expected_bboxs), err_msg)

        polygon = np.array([[91, 46], [144, 46], [144, 111], [91, 111]])
        transformed_polygons = transforms.apply_polygons([polygon])
        expected_polygon = np.array([[934.0, 33.0], [934.0, 80.0], [896.0, 80.0], [896.0, 33.0]])
        self.assertEqual(1, len(transformed_polygons))
        err_msg = "transformed_polygon = {}, expected {}".format(
            transformed_polygons[0], expected_polygon
        )
        self.assertTrue(polygon_allclose(transformed_polygons[0], expected_polygon), err_msg)

    def test_apply_rotated_boxes_unequal_scaling_factor(self):
        np.random.seed(125)
        h, w = 400, 200
        newh, neww = 800, 800
        image = np.random.rand(h, w)
        augs = []
        augs.append(T.Resize(shape=(newh, neww)))
        image, transforms = T.apply_augmentations(augs, image)
        image_shape = image.shape[:2]  # h, w
        assert image_shape == (newh, neww)

        boxes = np.array(
            [
                [150, 100, 40, 20, 0],
                [150, 100, 40, 20, 30],
                [150, 100, 40, 20, 90],
                [150, 100, 40, 20, -90],
            ],
            dtype=np.float64,
        )
        transformed_boxes = transforms.apply_rotated_box(boxes)

        expected_bboxes = np.array(
            [
                [600, 200, 160, 40, 0],
                [600, 200, 144.22205102, 52.91502622, 49.10660535],
                [600, 200, 80, 80, 90],
                [600, 200, 80, 80, -90],
            ],
            dtype=np.float64,
        )
        err_msg = "transformed_boxes = {}, expected {}".format(transformed_boxes, expected_bboxes)
        assert np.allclose(transformed_boxes, expected_bboxes), err_msg

    def test_print_augmentation(self):
        t = T.RandomCrop("relative", (100, 100))
        self.assertEqual(str(t), "RandomCrop(crop_type='relative', crop_size=(100, 100))")

        t0 = T.RandomFlip(prob=0.5)
        self.assertEqual(str(t0), "RandomFlip(prob=0.5)")

        t1 = T.RandomFlip()
        self.assertEqual(str(t1), "RandomFlip()")

        t = T.AugmentationList([t0, t1])
        self.assertEqual(str(t), f"AugmentationList[{t0}, {t1}]")

    def test_random_apply_prob_out_of_range_check(self):
        test_probabilities = {0.0: True, 0.5: True, 1.0: True, -0.01: False, 1.01: False}

        for given_probability, is_valid in test_probabilities.items():
            if not is_valid:
                self.assertRaises(AssertionError, T.RandomApply, None, prob=given_probability)
            else:
                T.RandomApply(T.NoOpTransform(), prob=given_probability)

    def test_random_apply_wrapping_aug_probability_occured_evaluation(self):
        transform_mock = mock.MagicMock(name="MockTransform", spec=T.Augmentation)
        image_mock = mock.MagicMock(name="MockImage")
        random_apply = T.RandomApply(transform_mock, prob=0.001)

        with mock.patch.object(random_apply, "_rand_range", return_value=0.0001):
            transform = random_apply.get_transform(image_mock)
        transform_mock.get_transform.assert_called_once_with(image_mock)
        self.assertIsNot(transform, transform_mock)

    def test_random_apply_wrapping_std_transform_probability_occured_evaluation(self):
        transform_mock = mock.MagicMock(name="MockTransform", spec=T.Transform)
        image_mock = mock.MagicMock(name="MockImage")
        random_apply = T.RandomApply(transform_mock, prob=0.001)

        with mock.patch.object(random_apply, "_rand_range", return_value=0.0001):
            transform = random_apply.get_transform(image_mock)
        self.assertIs(transform, transform_mock)

    def test_random_apply_probability_not_occured_evaluation(self):
        transform_mock = mock.MagicMock(name="MockTransform", spec=T.Augmentation)
        image_mock = mock.MagicMock(name="MockImage")
        random_apply = T.RandomApply(transform_mock, prob=0.001)

        with mock.patch.object(random_apply, "_rand_range", return_value=0.9):
            transform = random_apply.get_transform(image_mock)
        transform_mock.get_transform.assert_not_called()
        self.assertIsInstance(transform, T.NoOpTransform)

    def test_augmentation_input_args(self):
        input_shape = (100, 100)
        output_shape = (50, 50)

        # define two augmentations with different args
        class TG1(T.Augmentation):
            def get_transform(self, image, sem_seg):
                return T.ResizeTransform(
                    input_shape[0], input_shape[1], output_shape[0], output_shape[1]
                )

        class TG2(T.Augmentation):
            def get_transform(self, image):
                assert image.shape[:2] == output_shape  # check that TG1 is applied
                return T.HFlipTransform(output_shape[1])

        image = np.random.rand(*input_shape).astype("float32")
        sem_seg = (np.random.rand(*input_shape) < 0.5).astype("uint8")
        inputs = T.AugInput(image, sem_seg=sem_seg)  # provide two args
        tfms = inputs.apply_augmentations([TG1(), TG2()])
        self.assertIsInstance(tfms[0], T.ResizeTransform)
        self.assertIsInstance(tfms[1], T.HFlipTransform)
        self.assertTrue(inputs.image.shape[:2] == output_shape)
        self.assertTrue(inputs.sem_seg.shape[:2] == output_shape)

        class TG3(T.Augmentation):
            def get_transform(self, image, nonexist):
                pass

        with self.assertRaises(AttributeError):
            inputs.apply_augmentations([TG3()])

    def test_augmentation_list(self):
        input_shape = (100, 100)
        image = np.random.rand(*input_shape).astype("float32")
        sem_seg = (np.random.rand(*input_shape) < 0.5).astype("uint8")
        inputs = T.AugInput(image, sem_seg=sem_seg)  # provide two args

        augs = T.AugmentationList([T.RandomFlip(), T.Resize(20)])
        _ = T.AugmentationList([augs, T.Resize(30)])(inputs)
        # 3 in latest fvcore (flattened transformlist), 2 in older
        # self.assertEqual(len(tfms), 3)

    def test_color_transforms(self):
        rand_img = np.random.random((100, 100, 3)) * 255
        rand_img = rand_img.astype("uint8")

        # Test no-op
        noop_transform = T.ColorTransform(lambda img: img)
        self.assertTrue(np.array_equal(rand_img, noop_transform.apply_image(rand_img)))

        # Test a ImageOps operation
        magnitude = np.random.randint(0, 256)
        solarize_transform = T.PILColorTransform(lambda img: ImageOps.solarize(img, magnitude))
        expected_img = ImageOps.solarize(Image.fromarray(rand_img), magnitude)
        self.assertTrue(np.array_equal(expected_img, solarize_transform.apply_image(rand_img)))

    def test_resize_transform(self):
        input_shapes = [(100, 100), (100, 100, 1), (100, 100, 3)]
        output_shapes = [(200, 200), (200, 200, 1), (200, 200, 3)]
        for in_shape, out_shape in zip(input_shapes, output_shapes):
            in_img = np.random.randint(0, 255, size=in_shape, dtype=np.uint8)
            tfm = T.ResizeTransform(in_shape[0], in_shape[1], out_shape[0], out_shape[1])
            out_img = tfm.apply_image(in_img)
            self.assertEqual(out_img.shape, out_shape)

    def test_resize_shorted_edge_scriptable(self):
        def f(image):
            newh, neww = T.ResizeShortestEdge.get_output_shape(
                image.shape[-2], image.shape[-1], 80, 133
            )
            return F.interpolate(image.unsqueeze(0), size=(newh, neww))

        input = torch.randn(3, 10, 10)
        script_f = torch.jit.script(f)
        self.assertTrue(torch.allclose(f(input), script_f(input)))

        # generalize to new shapes
        input = torch.randn(3, 8, 100)
        self.assertTrue(torch.allclose(f(input), script_f(input)))

    def test_extent_transform(self):
        input_shapes = [(100, 100), (100, 100, 1), (100, 100, 3)]
        src_rect = (20, 20, 80, 80)
        output_shapes = [(200, 200), (200, 200, 1), (200, 200, 3)]
        for in_shape, out_shape in zip(input_shapes, output_shapes):
            in_img = np.random.randint(0, 255, size=in_shape, dtype=np.uint8)
            tfm = T.ExtentTransform(src_rect, out_shape[:2])
            out_img = tfm.apply_image(in_img)
            self.assertTrue(out_img.shape == out_shape)