|
|
import os |
|
|
from unittest import TestCase |
|
|
|
|
|
import cv2 |
|
|
import numpy as np |
|
|
import torch |
|
|
from mmengine.structures import InstanceData, PixelData |
|
|
|
|
|
from mmdet.evaluation import INSTANCE_OFFSET |
|
|
from mmdet.structures import DetDataSample |
|
|
from mmdet.visualization import DetLocalVisualizer |
|
|
|
|
|
|
|
|
def _rand_bboxes(num_boxes, h, w): |
|
|
cx, cy, bw, bh = torch.rand(num_boxes, 4).T |
|
|
|
|
|
tl_x = ((cx * w) - (w * bw / 2)).clamp(0, w) |
|
|
tl_y = ((cy * h) - (h * bh / 2)).clamp(0, h) |
|
|
br_x = ((cx * w) + (w * bw / 2)).clamp(0, w) |
|
|
br_y = ((cy * h) + (h * bh / 2)).clamp(0, h) |
|
|
|
|
|
bboxes = torch.stack([tl_x, tl_y, br_x, br_y], dim=0).T |
|
|
return bboxes |
|
|
|
|
|
|
|
|
def _create_panoptic_data(num_boxes, h, w): |
|
|
sem_seg = np.zeros((h, w), dtype=np.int64) + 2 |
|
|
bboxes = _rand_bboxes(num_boxes, h, w).int() |
|
|
labels = torch.randint(2, (num_boxes, )) |
|
|
for i in range(num_boxes): |
|
|
x, y, w, h = bboxes[i] |
|
|
sem_seg[y:y + h, x:x + w] = (i + 1) * INSTANCE_OFFSET + labels[i] |
|
|
|
|
|
return sem_seg[None] |
|
|
|
|
|
|
|
|
class TestDetLocalVisualizer(TestCase): |
|
|
|
|
|
def test_add_datasample(self): |
|
|
h = 12 |
|
|
w = 10 |
|
|
num_class = 3 |
|
|
num_bboxes = 5 |
|
|
out_file = 'out_file.jpg' |
|
|
|
|
|
image = np.random.randint(0, 256, size=(h, w, 3)).astype('uint8') |
|
|
|
|
|
|
|
|
gt_instances = InstanceData() |
|
|
gt_instances.bboxes = _rand_bboxes(num_bboxes, h, w) |
|
|
gt_instances.labels = torch.randint(0, num_class, (num_bboxes, )) |
|
|
det_data_sample = DetDataSample() |
|
|
det_data_sample.gt_instances = gt_instances |
|
|
|
|
|
det_local_visualizer = DetLocalVisualizer() |
|
|
det_local_visualizer.add_datasample( |
|
|
'image', image, det_data_sample, draw_pred=False) |
|
|
|
|
|
|
|
|
det_local_visualizer.add_datasample( |
|
|
'image', |
|
|
image, |
|
|
det_data_sample, |
|
|
draw_pred=False, |
|
|
out_file=out_file) |
|
|
assert os.path.exists(out_file) |
|
|
drawn_img = cv2.imread(out_file) |
|
|
assert drawn_img.shape == (h, w, 3) |
|
|
os.remove(out_file) |
|
|
|
|
|
|
|
|
pred_instances = InstanceData() |
|
|
pred_instances.bboxes = _rand_bboxes(num_bboxes, h, w) |
|
|
pred_instances.labels = torch.randint(0, num_class, (num_bboxes, )) |
|
|
pred_instances.scores = torch.rand((num_bboxes, )) |
|
|
det_data_sample.pred_instances = pred_instances |
|
|
|
|
|
det_local_visualizer.add_datasample( |
|
|
'image', image, det_data_sample, out_file=out_file) |
|
|
self._assert_image_and_shape(out_file, (h, w * 2, 3)) |
|
|
|
|
|
det_local_visualizer.add_datasample( |
|
|
'image', image, det_data_sample, draw_gt=False, out_file=out_file) |
|
|
self._assert_image_and_shape(out_file, (h, w, 3)) |
|
|
|
|
|
det_local_visualizer.add_datasample( |
|
|
'image', |
|
|
image, |
|
|
det_data_sample, |
|
|
draw_pred=False, |
|
|
out_file=out_file) |
|
|
self._assert_image_and_shape(out_file, (h, w, 3)) |
|
|
|
|
|
|
|
|
det_local_visualizer.dataset_meta = dict(classes=('1', '2')) |
|
|
gt_sem_seg = _create_panoptic_data(num_bboxes, h, w) |
|
|
panoptic_seg = PixelData(sem_seg=gt_sem_seg) |
|
|
|
|
|
det_data_sample = DetDataSample() |
|
|
det_data_sample.gt_panoptic_seg = panoptic_seg |
|
|
|
|
|
pred_sem_seg = _create_panoptic_data(num_bboxes, h, w) |
|
|
panoptic_seg = PixelData(sem_seg=pred_sem_seg) |
|
|
det_data_sample.pred_panoptic_seg = panoptic_seg |
|
|
|
|
|
det_local_visualizer.add_datasample( |
|
|
'image', image, det_data_sample, out_file=out_file) |
|
|
self._assert_image_and_shape(out_file, (h, w * 2, 3)) |
|
|
|
|
|
|
|
|
det_local_visualizer.dataset_meta = {} |
|
|
with self.assertRaises(AssertionError): |
|
|
det_local_visualizer.add_datasample( |
|
|
'image', image, det_data_sample, out_file=out_file) |
|
|
|
|
|
def _assert_image_and_shape(self, out_file, out_shape): |
|
|
assert os.path.exists(out_file) |
|
|
drawn_img = cv2.imread(out_file) |
|
|
assert drawn_img.shape == out_shape |
|
|
os.remove(out_file) |
|
|
|
