mmocr/tests/test_visualization/test_kie_visualizer.py

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

import cv2
import numpy as np
import torch
from mmengine.structures import InstanceData

from mmocr.structures import KIEDataSample
from mmocr.utils import bbox2poly
from mmocr.visualization import KIELocalVisualizer


class TestTextKIELocalVisualizer(unittest.TestCase):

    def setUp(self):
        h, w = 12, 10
        self.image = np.random.randint(0, 256, size=(h, w, 3)).astype('uint8')
        edge_labels = torch.rand((5, 5)) > 0.5
        # gt_instances
        data_sample = KIEDataSample()
        gt_instances_data = dict(
            bboxes=self._rand_bboxes(5, h, w),
            polygons=self._rand_polys(5, h, w),
            labels=torch.zeros(5, ),
            texts=['text1', 'text2', 'text3', 'text4', 'text5'],
            edge_labels=edge_labels)
        gt_instances = InstanceData(**gt_instances_data)
        data_sample.gt_instances = gt_instances

        pred_instances_data = dict(
            bboxes=self._rand_bboxes(5, h, w),
            labels=torch.zeros(5, ),
            scores=torch.rand((5, )),
            texts=['text1', 'text2', 'text3', 'text4', 'text5'],
            edge_labels=edge_labels)
        pred_instances = InstanceData(**pred_instances_data)
        data_sample.pred_instances = pred_instances
        data_sample = data_sample.numpy()
        self.data_sample = data_sample

    @staticmethod
    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).unsqueeze(0)
        tl_y = ((cy * h) - (h * bh / 2)).clamp(0, h).unsqueeze(0)
        br_x = ((cx * w) + (w * bw / 2)).clamp(0, w).unsqueeze(0)
        br_y = ((cy * h) + (h * bh / 2)).clamp(0, h).unsqueeze(0)

        bboxes = torch.cat([tl_x, tl_y, br_x, br_y], dim=0).T

        return bboxes

    def _rand_polys(self, num_bboxes, h, w):
        bboxes = self._rand_bboxes(num_bboxes, h, w)
        bboxes = bboxes.tolist()
        polys = [bbox2poly(bbox) for bbox in bboxes]
        return polys

    def test_add_datasample(self):
        image = self.image
        h, w, c = image.shape

        visualizer = KIELocalVisualizer(is_openset=True)
        visualizer.dataset_meta = dict(category=[
            dict(id=0, name='bg'),
            dict(id=1, name='key'),
            dict(id=2, name='value'),
            dict(id=3, name='other')
        ])
        visualizer.add_datasample('image', image, self.data_sample)

        with tempfile.TemporaryDirectory() as tmp_dir:
            # test out
            out_file = osp.join(tmp_dir, 'out_file.jpg')
            visualizer.add_datasample(
                'image',
                image,
                self.data_sample,
                out_file=out_file,
                draw_gt=False,
                draw_pred=False)
            self._assert_image_and_shape(out_file, (h, w, c))

            visualizer.add_datasample(
                'image', image, self.data_sample, out_file=out_file)
            self._assert_image_and_shape(out_file, (h * 2, w * 4, c))

            visualizer.add_datasample(
                'image',
                image,
                self.data_sample,
                draw_gt=False,
                out_file=out_file)
            self._assert_image_and_shape(out_file, (h, w * 4, c))

            visualizer.add_datasample(
                'image',
                image,
                self.data_sample,
                draw_pred=False,
                out_file=out_file)
            self._assert_image_and_shape(out_file, (h, w * 4, c))

    def _assert_image_and_shape(self, out_file, out_shape):
        self.assertTrue(osp.exists(out_file))
        drawn_img = cv2.imread(out_file)
        self.assertTrue(drawn_img.shape == out_shape)