Spaces:
Running
Running
File size: 4,385 Bytes
174ad5e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 |
# 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))
visualizer = KIELocalVisualizer(is_openset=False)
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,
draw_pred=False,
out_file=out_file)
self._assert_image_and_shape(out_file, (h, w * 3, 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)
|