mmdetection/tests/test_datasets/test_transforms/test_formatting.py

# Copyright (c) OpenMMLab. All rights reserved.
import copy
import os.path as osp
import unittest

import numpy as np
import torch
from mmengine.structures import InstanceData, LabelData, PixelData

from mmdet.datasets.transforms import (PackDetInputs, PackReIDInputs,
                                       PackTrackInputs)
from mmdet.structures import DetDataSample, ReIDDataSample
from mmdet.structures.mask import BitmapMasks


class TestPackDetInputs(unittest.TestCase):

    def setUp(self):
        """Setup the model and optimizer which are used in every test method.

        TestCase calls functions in this order: setUp() -> testMethod() ->
        tearDown() -> cleanUp()
        """
        data_prefix = osp.join(osp.dirname(__file__), '../../data')
        img_path = osp.join(data_prefix, 'color.jpg')
        rng = np.random.RandomState(0)
        self.results1 = {
            'img_id': 1,
            'img_path': img_path,
            'ori_shape': (300, 400),
            'img_shape': (600, 800),
            'scale_factor': 2.0,
            'flip': False,
            'img': rng.rand(300, 400),
            'gt_seg_map': rng.rand(300, 400),
            'gt_masks':
            BitmapMasks(rng.rand(3, 300, 400), height=300, width=400),
            'gt_bboxes_labels': rng.rand(3, ),
            'gt_ignore_flags': np.array([0, 0, 1], dtype=bool),
            'proposals': rng.rand(2, 4),
            'proposals_scores': rng.rand(2, )
        }
        self.results2 = {
            'img_id': 1,
            'img_path': img_path,
            'ori_shape': (300, 400),
            'img_shape': (600, 800),
            'scale_factor': 2.0,
            'flip': False,
            'img': rng.rand(300, 400),
            'gt_seg_map': rng.rand(300, 400),
            'gt_masks':
            BitmapMasks(rng.rand(3, 300, 400), height=300, width=400),
            'gt_bboxes_labels': rng.rand(3, ),
            'proposals': rng.rand(2, 4),
            'proposals_scores': rng.rand(2, )
        }
        self.meta_keys = ('img_id', 'img_path', 'ori_shape', 'scale_factor',
                          'flip')

    def test_transform(self):
        transform = PackDetInputs(meta_keys=self.meta_keys)
        results = transform(copy.deepcopy(self.results1))
        self.assertIn('data_samples', results)
        self.assertIsInstance(results['data_samples'], DetDataSample)
        self.assertIsInstance(results['data_samples'].gt_instances,
                              InstanceData)
        self.assertIsInstance(results['data_samples'].ignored_instances,
                              InstanceData)
        self.assertEqual(len(results['data_samples'].gt_instances), 2)
        self.assertEqual(len(results['data_samples'].ignored_instances), 1)
        self.assertIsInstance(results['data_samples'].gt_sem_seg, PixelData)
        self.assertIsInstance(results['data_samples'].proposals, InstanceData)
        self.assertEqual(len(results['data_samples'].proposals), 2)
        self.assertIsInstance(results['data_samples'].proposals.bboxes,
                              torch.Tensor)
        self.assertIsInstance(results['data_samples'].proposals.scores,
                              torch.Tensor)

    def test_transform_without_ignore(self):
        transform = PackDetInputs(meta_keys=self.meta_keys)
        results = transform(copy.deepcopy(self.results2))
        self.assertIn('data_samples', results)
        self.assertIsInstance(results['data_samples'], DetDataSample)
        self.assertIsInstance(results['data_samples'].gt_instances,
                              InstanceData)
        self.assertIsInstance(results['data_samples'].ignored_instances,
                              InstanceData)
        self.assertEqual(len(results['data_samples'].gt_instances), 3)
        self.assertEqual(len(results['data_samples'].ignored_instances), 0)
        self.assertIsInstance(results['data_samples'].gt_sem_seg, PixelData)
        self.assertIsInstance(results['data_samples'].proposals, InstanceData)
        self.assertEqual(len(results['data_samples'].proposals), 2)
        self.assertIsInstance(results['data_samples'].proposals.bboxes,
                              torch.Tensor)
        self.assertIsInstance(results['data_samples'].proposals.scores,
                              torch.Tensor)

    def test_repr(self):
        transform = PackDetInputs(meta_keys=self.meta_keys)
        self.assertEqual(
            repr(transform), f'PackDetInputs(meta_keys={self.meta_keys})')


class TestPackTrackInputs(unittest.TestCase):

    def setUp(self):
        self.H, self.W = 5, 10
        self.img = np.zeros((self.H, self.W, 3))
        self.gt_bboxes = np.zeros((2, 4))
        self.gt_masks = BitmapMasks(
            np.random.rand(2, self.H, self.W), height=self.H, width=self.W)
        self.gt_bboxes_labels = [
            np.zeros((2, )),
            np.zeros((2, )) + 1,
            np.zeros((2, )) - 1
        ]
        self.gt_instances_ids = [
            np.ones((2, ), dtype=np.int32),
            np.ones((2, ), dtype=np.int32) - 1,
            np.ones((2, ), dtype=np.int32) + 1
        ]
        self.frame_id = [0, 1, 2]
        self.scale_factor = [1.0, 1.5, 2.0]
        self.flip = [False] * 3
        self.ori_shape = [(self.H, self.W)] * 3
        self.img_id = [0, 1, 2]
        self.results_1 = dict(
            img=[self.img.copy(),
                 self.img.copy(),
                 self.img.copy()],
            gt_bboxes=[
                self.gt_bboxes.copy(),
                self.gt_bboxes.copy(),
                self.gt_bboxes.copy()
            ],
            gt_bboxes_labels=copy.deepcopy(self.gt_bboxes_labels),
            gt_instances_ids=copy.deepcopy(self.gt_instances_ids),
            gt_masks=[
                copy.deepcopy(self.gt_masks),
                copy.deepcopy(self.gt_masks),
                copy.deepcopy(self.gt_masks)
            ],
            frame_id=self.frame_id,
            ori_shape=self.ori_shape,
            scale_factor=self.scale_factor,
            flip=self.flip,
            img_id=self.img_id,
            key_frame_flags=[False, True, False])

        self.results_2 = copy.deepcopy(self.results_1)
        self.gt_ignore_flags = [
            np.array([0, 1], dtype=np.bool),
            np.array([1, 0], dtype=np.bool),
            np.array([0, 0], dtype=np.bool)
        ]
        self.results_2.update(
            dict(gt_ignore_flags=copy.deepcopy(self.gt_ignore_flags)))

        self.meta_keys = ('frame_id', 'ori_shape', 'scale_factor', 'flip')
        self.pack_track_inputs = PackTrackInputs(meta_keys=self.meta_keys)

    def test_transform_without_ignore(self):
        track_results = self.pack_track_inputs(self.results_1)
        assert isinstance(track_results, dict)

        inputs = track_results['inputs']
        assert isinstance(inputs, torch.Tensor)
        assert inputs.shape == (3, 3, self.H, self.W)

        track_data_sample = track_results['data_samples']
        assert len(track_data_sample) == 3
        assert 'key_frames_inds' in track_data_sample.metainfo and \
            track_data_sample.key_frames_inds == [1]
        assert 'ref_frames_inds' in track_data_sample.metainfo and \
            track_data_sample.ref_frames_inds == [0, 2]
        for i, data_sample in enumerate(track_data_sample):
            assert data_sample.gt_instances.bboxes.shape == (2, 4)
            assert len(data_sample.gt_instances.masks) == 2
            assert (data_sample.gt_instances.labels.numpy() ==
                    self.gt_bboxes_labels[i]).all()
            assert (data_sample.gt_instances.instances_ids.numpy() ==
                    self.gt_instances_ids[i]).all()
            for key in self.meta_keys:
                assert data_sample.metainfo[key] == getattr(self, key)[i]

    def test_transform_with_ignore(self):
        track_results = self.pack_track_inputs(self.results_2)
        assert isinstance(track_results, dict)

        inputs = track_results['inputs']
        assert isinstance(inputs, torch.Tensor)
        assert inputs.shape == (3, 3, self.H, self.W)

        track_data_sample = track_results['data_samples']
        assert len(track_data_sample) == 3
        for i, data_sample in enumerate(track_data_sample):
            valid_mask = ~self.gt_ignore_flags[i]
            valid_len = valid_mask.sum().item()
            assert data_sample.gt_instances.bboxes.shape == (valid_len, 4)
            assert len(data_sample.gt_instances.masks) == valid_len
            assert (data_sample.gt_instances.labels.numpy() ==
                    self.gt_bboxes_labels[i][valid_mask]).all()
            assert (data_sample.gt_instances.instances_ids.numpy() ==
                    self.gt_instances_ids[i][valid_mask]).all()
            for key in self.meta_keys:
                assert data_sample.metainfo[key] == getattr(self, key)[i]


class TestPackReIDInputs(unittest.TestCase):

    @classmethod
    def setUpClass(cls):
        cls.results = dict(
            img=np.random.randn(256, 128, 3),
            gt_label=0,
            img_path='',
            ori_shape=(128, 128),
            img_shape=(256, 128),
            scale=(128, 256),
            scale_factor=(1., 2.),
            flip=False,
            flip_direction=None)
        cls.pack_reid_inputs = PackReIDInputs(
            meta_keys=('flip', 'flip_direction'))

    def test_transform(self):
        results = self.pack_reid_inputs(self.results)
        self.assertIn('inputs', results)
        self.assertIsInstance(results['inputs'], torch.Tensor)
        self.assertIn('data_samples', results)
        data_sample = results['data_samples']
        self.assertIsInstance(data_sample, ReIDDataSample)
        self.assertIsInstance(data_sample.gt_label, LabelData)
        self.assertEqual(data_sample.img_path, '')
        self.assertEqual(data_sample.ori_shape, (128, 128))
        self.assertEqual(data_sample.img_shape, (256, 128))
        self.assertEqual(data_sample.scale, (128, 256))
        self.assertEqual(data_sample.scale_factor, (1., 2.))
        self.assertEqual(data_sample.flip, False)
        self.assertIsNone(data_sample.flip_direction)

    def test_repr(self):
        self.assertEqual(
            repr(self.pack_reid_inputs),
            f'PackReIDInputs(meta_keys={self.pack_reid_inputs.meta_keys})')