mirror of https://github.com/open-mmlab/mmocr.git
56 lines
2.2 KiB
Python
56 lines
2.2 KiB
Python
# Copyright (c) OpenMMLab. All rights reserved.
|
|
from unittest import TestCase
|
|
|
|
import numpy as np
|
|
|
|
from mmocr.visualization import BaseLocalVisualizer
|
|
|
|
|
|
class TestBaseLocalVisualizer(TestCase):
|
|
|
|
def test_get_labels_image(self):
|
|
labels = ['a', 'b', 'c']
|
|
image = np.zeros((40, 40, 3), dtype=np.uint8)
|
|
bboxes = np.array([[0, 0, 10, 10], [10, 10, 20, 20], [20, 20, 30, 30]])
|
|
labels_image = BaseLocalVisualizer().get_labels_image(
|
|
image,
|
|
labels,
|
|
bboxes=bboxes,
|
|
auto_font_size=True,
|
|
colors=['r', 'r', 'r', 'r'])
|
|
self.assertEqual(labels_image.shape, (40, 40, 3))
|
|
|
|
def test_get_polygons_image(self):
|
|
polygons = [np.array([0, 0, 10, 10, 20, 20, 30, 30]).reshape(-1, 2)]
|
|
image = np.zeros((40, 40, 3), dtype=np.uint8)
|
|
polygons_image = BaseLocalVisualizer().get_polygons_image(
|
|
image, polygons, colors=['r', 'r', 'r', 'r'])
|
|
self.assertEqual(polygons_image.shape, (40, 40, 3))
|
|
|
|
polygons_image = BaseLocalVisualizer().get_polygons_image(
|
|
image, polygons, colors=['r', 'r', 'r', 'r'], filling=True)
|
|
self.assertEqual(polygons_image.shape, (40, 40, 3))
|
|
|
|
def test_get_bboxes_image(self):
|
|
bboxes = np.array([[0, 0, 10, 10], [10, 10, 20, 20], [20, 20, 30, 30]])
|
|
image = np.zeros((40, 40, 3), dtype=np.uint8)
|
|
bboxes_image = BaseLocalVisualizer().get_bboxes_image(
|
|
image, bboxes, colors=['r', 'r', 'r', 'r'])
|
|
self.assertEqual(bboxes_image.shape, (40, 40, 3))
|
|
|
|
bboxes_image = BaseLocalVisualizer().get_bboxes_image(
|
|
image, bboxes, colors=['r', 'r', 'r', 'r'], filling=True)
|
|
self.assertEqual(bboxes_image.shape, (40, 40, 3))
|
|
|
|
def test_cat_images(self):
|
|
image1 = np.zeros((40, 40, 3), dtype=np.uint8)
|
|
image2 = np.zeros((40, 40, 3), dtype=np.uint8)
|
|
image = BaseLocalVisualizer()._cat_image([image1, image2], axis=1)
|
|
self.assertEqual(image.shape, (40, 80, 3))
|
|
|
|
image = BaseLocalVisualizer()._cat_image([], axis=0)
|
|
self.assertIsNone(image)
|
|
|
|
image = BaseLocalVisualizer()._cat_image([image1, None], axis=0)
|
|
self.assertEqual(image.shape, (40, 40, 3))
|