mmdetection/tests/test_visualization/test_local_visualizer.py

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, TrackLocalVisualizer


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')

        # test gt_instances
        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)

        # test out_file
        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)

        # test gt_instances and pred_instances
        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))

        # test gt_panoptic_seg and pred_panoptic_seg
        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))

        # class information must be provided
        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)


class TestTrackLocalVisualizer(TestCase):

    @staticmethod
    def _get_gt_instances():
        bboxes = np.array([[912, 484, 1009, 593], [1338, 418, 1505, 797]])
        masks = np.zeros((2, 1080, 1920), dtype=np.bool_)
        for i, bbox in enumerate(bboxes):
            masks[i, bbox[1]:bbox[3], bbox[0]:bbox[2]] = True
        instances_data = dict(
            bboxes=torch.tensor(bboxes),
            masks=masks,
            instances_id=torch.tensor([1, 2]),
            labels=torch.tensor([0, 1]))
        instances = InstanceData(**instances_data)
        return instances

    @staticmethod
    def _get_pred_instances():
        instances_data = dict(
            bboxes=torch.tensor([[900, 500, 1000, 600], [1300, 400, 1500,
                                                         800]]),
            instances_id=torch.tensor([1, 2]),
            labels=torch.tensor([0, 1]),
            scores=torch.tensor([0.955, 0.876]))
        instances = InstanceData(**instances_data)
        return instances

    @staticmethod
    def _assert_image_and_shape(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)

    def test_add_datasample(self):
        out_file = 'out_file.jpg'
        h, w = 1080, 1920
        image = np.random.randint(0, 256, size=(h, w, 3)).astype('uint8')
        gt_instances = self._get_gt_instances()
        pred_instances = self._get_pred_instances()
        image_data_sample = DetDataSample()
        image_data_sample.gt_instances = gt_instances
        image_data_sample.pred_track_instances = pred_instances

        track_local_visualizer = TrackLocalVisualizer(alpha=0.2)
        track_local_visualizer.dataset_meta = dict(
            classes=['pedestrian', 'vehicle'])

        # test gt_instances
        track_local_visualizer.add_datasample('image', image,
                                              image_data_sample, None)

        # test out_file
        track_local_visualizer.add_datasample(
            'image', image, image_data_sample, None, out_file=out_file)
        self._assert_image_and_shape(out_file, (h, w, 3))

        # test gt_instances and pred_instances
        track_local_visualizer.add_datasample(
            'image', image, image_data_sample, out_file=out_file)
        self._assert_image_and_shape(out_file, (h, 2 * w, 3))

        track_local_visualizer.add_datasample(
            'image',
            image,
            image_data_sample,
            draw_gt=False,
            out_file=out_file)
        self._assert_image_and_shape(out_file, (h, w, 3))

        track_local_visualizer.add_datasample(
            'image',
            image,
            image_data_sample,
            draw_pred=False,
            out_file=out_file)
        self._assert_image_and_shape(out_file, (h, w, 3))