test_data_transforms_crop.py

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


import unittest

import d2go.data.transforms.box_utils as bu
import numpy as np
import torch
from d2go.data.transforms import crop as tf_crop


class TestDataTransformsCrop(unittest.TestCase):
    def test_transform_crop_extent_transform(self):
        img_wh = (16, 11)

        sem_seg = np.zeros([img_wh[1], img_wh[0]], dtype=np.uint8)
        # h, w
        sem_seg[5, 4] = 1
        sem_seg[10, 13] = 1
        sem_seg[5:11, 4:14] = 1

        # src_rect: [x0, y0, x1, y1] in pixel coordinate, output_size: [h, w]
        trans = tf_crop.ExtentTransform(src_rect=[4, 5, 14, 11], output_size=[6, 10])
        out_mask = trans.apply_segmentation(sem_seg)
        self.assertArrayEqual(out_mask.shape, torch.Tensor([6, 10]))
        self.assertArrayEqual(np.unique(out_mask), torch.Tensor([1]))

        trans = tf_crop.ExtentTransform(src_rect=[3, 4, 15, 11], output_size=[7, 12])
        out_mask = trans.apply_segmentation(sem_seg)
        self.assertArrayEqual(out_mask.shape, torch.Tensor([7, 12]))
        self.assertArrayEqual(np.unique(out_mask), torch.Tensor([0, 1]))
        self.assertArrayEqual(np.unique(out_mask[1:, 1:-1]), torch.Tensor([1]))
        self.assertEqual(out_mask[:, 0].sum(), 0)
        self.assertArrayEqual(out_mask[0, :].sum(), 0)
        self.assertArrayEqual(out_mask[:, -1].sum(), 0)

    def test_transform_crop_random_crop_fixed_aspect_ratio(self):
        aug = tf_crop.RandomCropFixedAspectRatio([1.0 / 2])
        img_wh = (16, 11)

        img = np.ones([img_wh[1], img_wh[0], 3], dtype=np.uint8)
        sem_seg = np.zeros([img_wh[1], img_wh[0]], dtype=np.uint8)
        sem_seg[5, 4] = 1
        sem_seg[10, 13] = 1
        mask_xywh = bu.get_box_from_mask(sem_seg)
        self.assertArrayEqual(mask_xywh, torch.Tensor([4, 5, 10, 6]))

        trans = aug.get_transform(img, sem_seg)
        self.assertArrayEqual(trans.src_rect, torch.Tensor([4, -2, 14, 18]))
        self.assertArrayEqual(trans.output_size, torch.Tensor([20, 10]))

        out_img = trans.apply_image(img)
        self.assertArrayEqual(out_img.shape, torch.Tensor([20, 10, 3]))

        self.assertArrayEqual(np.unique(out_img[2:13, :, :]), torch.Tensor([1]))
        self.assertArrayEqual(np.unique(out_img[0:2, :, :]), torch.Tensor([0]))
        self.assertArrayEqual(np.unique(out_img[13:, :, :]), torch.Tensor([0]))

        out_mask = trans.apply_segmentation(sem_seg)
        self.assertArrayEqual(out_mask.shape, torch.Tensor([20, 10]))
        self.assertEqual(out_mask[7, 0], 1)
        self.assertEqual(out_mask[12, -1], 1)

    def test_transform_crop_random_crop_fixed_aspect_ratio_scale_offset(self):
        aug = tf_crop.RandomCropFixedAspectRatio(
            [1.0 / 2], scale_range=[0.5, 0.5], offset_scale_range=[-0.5, -0.5]
        )
        img_wh = (16, 11)

        img = np.ones([img_wh[1], img_wh[0], 3], dtype=np.uint8)
        sem_seg = np.zeros([img_wh[1], img_wh[0]], dtype=np.uint8)
        sem_seg[5, 4] = 1
        sem_seg[10, 13] = 1
        sem_seg[5:11, 4:14] = 1
        mask_xywh = bu.get_box_from_mask(sem_seg)
        self.assertArrayEqual(mask_xywh, torch.Tensor([4, 5, 10, 6]))

        trans = aug.get_transform(img, sem_seg)
        self.assertArrayEqual(trans.src_rect, torch.Tensor([1.5, 0.0, 6.5, 10.0]))
        self.assertArrayEqual(trans.output_size, torch.Tensor([10, 5]))

        out_img = trans.apply_image(img)
        self.assertArrayEqual(out_img.shape, torch.Tensor([10, 5, 3]))
        self.assertEqual(np.unique(out_img), 1)

        out_mask = trans.apply_segmentation(sem_seg)
        self.assertArrayEqual(out_mask.shape, torch.Tensor([10, 5]))
        self.assertEqual(np.unique(out_mask[6:, 3:]), 1)

    def test_transform_crop_random_crop_fixed_aspect_ratio_empty_mask(self):
        """The sem_mask is empty (the whole image is background)"""
        aug = tf_crop.RandomCropFixedAspectRatio([1.0 / 2])
        img_wh = (16, 11)

        img = np.ones([img_wh[1], img_wh[0], 3], dtype=np.uint8)
        sem_seg = np.zeros([img_wh[1], img_wh[0]], dtype=np.uint8)

        mask_xywh = bu.get_box_from_mask(sem_seg)
        self.assertEqual(mask_xywh, None)

        trans = aug.get_transform(img, sem_seg)
        self.assertIsInstance(trans, tf_crop.NoOpTransform)

        out_img = trans.apply_image(img)
        self.assertArrayEqual(out_img.shape, img.shape)

        out_mask = trans.apply_segmentation(sem_seg)
        self.assertArrayEqual(out_mask.shape, sem_seg.shape)

    def test_pad_transform(self):
        crop_w, crop_h = 4, 3
        full_w, full_h = 11, 9
        crop_x, crop_y = 5, 6
        trans = tf_crop.PadTransform(crop_x, crop_y, crop_w, crop_h, full_w, full_h)
        img = np.ones([crop_h, crop_w])
        trans_img = trans.apply_image(img)
        self.assertArrayEqual(trans_img.shape, [full_h, full_w])
        self.assertArrayEqual(np.unique(trans_img), [0, 1])
        full_img_gt = np.zeros([full_h, full_w])
        full_img_gt[crop_y : (crop_y + crop_h), crop_x : (crop_x + crop_w)] = 1
        self.assertArrayEqual(full_img_gt, trans_img)

    def test_crop_transform_inverse(self):
        crop_w, crop_h = 4, 3
        full_w, full_h = 11, 9
        crop_x, crop_y = 5, 6
        trans = tf_crop.InvertibleCropTransform(
            crop_x, crop_y, crop_w, crop_h, full_w, full_h
        )

        full_img_gt = np.zeros([full_h, full_w])
        full_img_gt[crop_y : (crop_y + crop_h), crop_x : (crop_x + crop_w)] = 1
        crop_img_gt = np.ones([crop_h, crop_w])

        self.assertArrayEqual(trans.apply_image(full_img_gt), crop_img_gt)
        self.assertArrayEqual(trans.inverse().apply_image(crop_img_gt), full_img_gt)
        self.assertArrayEqual(
            trans.inverse().inverse().apply_image(full_img_gt), crop_img_gt
        )

    def test_pad_transform(self):
        img_h, img_w = 10, 7
        divisibility = 8

        aug = tf_crop.PadBorderDivisible(divisibility)

        img = np.ones([img_h, img_w, 3]) * 3
        trans = aug.get_transform(img)
        pad_img = trans.apply_image(img)
        self.assertEqual(pad_img.shape, (16, 8, 3))
        inverse_img = trans.inverse().apply_image(pad_img)
        self.assertEqual(inverse_img.shape, (10, 7, 3))
        self.assertArrayEqual(img, inverse_img)

        mask = np.ones([img_h, img_w]) * 2
        pad_mask = trans.apply_segmentation(mask)
        self.assertEqual(pad_mask.shape, (16, 8))
        inverse_mask = trans.inverse().apply_segmentation(pad_mask)
        self.assertEqual(inverse_mask.shape, (10, 7))
        self.assertArrayEqual(mask, inverse_mask)

    def test_random_instance_crop(self):
        from detectron2.data import detection_utils as du
        from detectron2.data.transforms.augmentation import (
            AugInput,
            AugmentationList,
        )
        from detectron2.structures import BoxMode

        aug = tf_crop.RandomInstanceCrop([1.0, 1.0])

        img_w, img_h = 10, 7
        annotations = [
            {
                "category_id": 0,
                "bbox": [1, 1, 4, 3],
                "bbox_mode": BoxMode.XYWH_ABS,
            },
            {
                "category_id": 0,
                "bbox": [2, 2, 4, 3],
                "bbox_mode": BoxMode.XYWH_ABS,
            },
            {
                "category_id": 0,
                "bbox": [6, 5, 3, 2],
                "bbox_mode": BoxMode.XYWH_ABS,
            },
        ]

        img = np.ones([img_h, img_w, 3]) * 3

        inputs = AugInput(image=img)
        # pass additional arguments
        inputs.annotations = annotations
        transforms = AugmentationList([aug])(inputs)

        self.assertIn(
            inputs.image.shape, [torch.Size([3, 4, 3]), torch.Size([2, 3, 3])]
        )

        # from dataset mapper unused annotations will be filtered out due to the
        #  iscrowd flag
        image_shape = inputs.image.shape[:2]
        annos = [
            du.transform_instance_annotations(
                obj,
                transforms,
                image_shape,
            )
            for obj in annotations
            if obj.get("iscrowd", 0) == 0
        ]
        instances = du.annotations_to_instances(annos, image_shape)
        filtered_instances = du.filter_empty_instances(instances)
        self.assertEqual(len(filtered_instances), 1)
        self.assertArrayEqual(
            filtered_instances.gt_boxes.tensor.tolist(),
            [[0, 0, image_shape[1], image_shape[0]]],
        )

    def assertArrayEqual(self, a1, a2):
        self.assertTrue(np.array_equal(a1, a2))