Spaces:
Running
on
T4
Running
on
T4
| # Copyright (c) OpenMMLab. All rights reserved. | |
| import unittest | |
| import numpy as np | |
| import torch | |
| from mmdet.structures import DetDataSample | |
| from mmdet.structures.bbox import HorizontalBoxes | |
| from mmengine.structures import InstanceData | |
| from mmyolo.datasets import BatchShapePolicy, yolov5_collate | |
| def _rand_bboxes(rng, num_boxes, w, h): | |
| cx, cy, bw, bh = rng.rand(num_boxes, 4).T | |
| tl_x = ((cx * w) - (w * bw / 2)).clip(0, w) | |
| tl_y = ((cy * h) - (h * bh / 2)).clip(0, h) | |
| br_x = ((cx * w) + (w * bw / 2)).clip(0, w) | |
| br_y = ((cy * h) + (h * bh / 2)).clip(0, h) | |
| bboxes = np.vstack([tl_x, tl_y, br_x, br_y]).T | |
| return bboxes | |
| class TestYOLOv5Collate(unittest.TestCase): | |
| def test_yolov5_collate(self): | |
| rng = np.random.RandomState(0) | |
| inputs = torch.randn((3, 10, 10)) | |
| data_samples = DetDataSample() | |
| gt_instances = InstanceData() | |
| bboxes = _rand_bboxes(rng, 4, 6, 8) | |
| gt_instances.bboxes = HorizontalBoxes(bboxes, dtype=torch.float32) | |
| labels = rng.randint(1, 2, size=len(bboxes)) | |
| gt_instances.labels = torch.LongTensor(labels) | |
| data_samples.gt_instances = gt_instances | |
| out = yolov5_collate([dict(inputs=inputs, data_samples=data_samples)]) | |
| self.assertIsInstance(out, dict) | |
| self.assertTrue(out['inputs'].shape == (1, 3, 10, 10)) | |
| self.assertTrue(out['data_samples'], dict) | |
| self.assertTrue(out['data_samples']['bboxes_labels'].shape == (4, 6)) | |
| out = yolov5_collate([dict(inputs=inputs, data_samples=data_samples)] * | |
| 2) | |
| self.assertIsInstance(out, dict) | |
| self.assertTrue(out['inputs'].shape == (2, 3, 10, 10)) | |
| self.assertTrue(out['data_samples'], dict) | |
| self.assertTrue(out['data_samples']['bboxes_labels'].shape == (8, 6)) | |
| def test_yolov5_collate_with_multi_scale(self): | |
| rng = np.random.RandomState(0) | |
| inputs = torch.randn((3, 10, 10)) | |
| data_samples = DetDataSample() | |
| gt_instances = InstanceData() | |
| bboxes = _rand_bboxes(rng, 4, 6, 8) | |
| gt_instances.bboxes = HorizontalBoxes(bboxes, dtype=torch.float32) | |
| labels = rng.randint(1, 2, size=len(bboxes)) | |
| gt_instances.labels = torch.LongTensor(labels) | |
| data_samples.gt_instances = gt_instances | |
| out = yolov5_collate([dict(inputs=inputs, data_samples=data_samples)], | |
| use_ms_training=True) | |
| self.assertIsInstance(out, dict) | |
| self.assertTrue(out['inputs'][0].shape == (3, 10, 10)) | |
| self.assertTrue(out['data_samples'], dict) | |
| self.assertTrue(out['data_samples']['bboxes_labels'].shape == (4, 6)) | |
| self.assertIsInstance(out['inputs'], list) | |
| self.assertIsInstance(out['data_samples']['bboxes_labels'], | |
| torch.Tensor) | |
| out = yolov5_collate( | |
| [dict(inputs=inputs, data_samples=data_samples)] * 2, | |
| use_ms_training=True) | |
| self.assertIsInstance(out, dict) | |
| self.assertTrue(out['inputs'][0].shape == (3, 10, 10)) | |
| self.assertTrue(out['data_samples'], dict) | |
| self.assertTrue(out['data_samples']['bboxes_labels'].shape == (8, 6)) | |
| self.assertIsInstance(out['inputs'], list) | |
| self.assertIsInstance(out['data_samples']['bboxes_labels'], | |
| torch.Tensor) | |
| class TestBatchShapePolicy(unittest.TestCase): | |
| def test_batch_shape_policy(self): | |
| src_data_infos = [{ | |
| 'height': 20, | |
| 'width': 100, | |
| }, { | |
| 'height': 11, | |
| 'width': 100, | |
| }, { | |
| 'height': 21, | |
| 'width': 100, | |
| }, { | |
| 'height': 30, | |
| 'width': 100, | |
| }, { | |
| 'height': 10, | |
| 'width': 100, | |
| }] | |
| expected_data_infos = [{ | |
| 'height': 10, | |
| 'width': 100, | |
| 'batch_shape': np.array([96, 672]) | |
| }, { | |
| 'height': 11, | |
| 'width': 100, | |
| 'batch_shape': np.array([96, 672]) | |
| }, { | |
| 'height': 20, | |
| 'width': 100, | |
| 'batch_shape': np.array([160, 672]) | |
| }, { | |
| 'height': 21, | |
| 'width': 100, | |
| 'batch_shape': np.array([160, 672]) | |
| }, { | |
| 'height': 30, | |
| 'width': 100, | |
| 'batch_shape': np.array([224, 672]) | |
| }] | |
| batch_shapes_policy = BatchShapePolicy(batch_size=2) | |
| out_data_infos = batch_shapes_policy(src_data_infos) | |
| for i in range(5): | |
| self.assertEqual( | |
| (expected_data_infos[i]['height'], | |
| expected_data_infos[i]['width']), | |
| (out_data_infos[i]['height'], out_data_infos[i]['width'])) | |
| self.assertTrue( | |
| np.allclose(expected_data_infos[i]['batch_shape'], | |
| out_data_infos[i]['batch_shape'])) | |