Spaces:
Sleeping
Sleeping
File size: 3,885 Bytes
0b4516f |
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 |
# 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 TextDetDataSample
from mmocr.utils import bbox2poly
from mmocr.visualization import TextDetLocalVisualizer
class TestTextDetLocalVisualizer(unittest.TestCase):
def setUp(self):
h, w = 12, 10
self.image = np.random.randint(0, 256, size=(h, w, 3)).astype('uint8')
# gt_instances
data_sample = TextDetDataSample()
gt_instances_data = dict(
bboxes=self._rand_bboxes(5, h, w),
polygons=self._rand_polys(5, h, w),
labels=torch.zeros(5, ))
gt_instances = InstanceData(**gt_instances_data)
data_sample.gt_instances = gt_instances
pred_instances_data = dict(
bboxes=self._rand_bboxes(5, h, w),
polygons=self._rand_polys(5, h, w),
labels=torch.zeros(5, ),
scores=torch.rand((5, )))
pred_instances = InstanceData(**pred_instances_data)
data_sample.pred_instances = pred_instances
self.data_sample = data_sample
def test_text_det_local_visualizer(self):
for with_poly in [True, False]:
for with_bbox in [True, False]:
vis_cfg = dict(with_poly=with_poly, with_bbox=with_bbox)
self._test_add_datasample(vis_cfg=vis_cfg)
@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, vis_cfg):
image = self.image
h, w, c = image.shape
det_local_visualizer = TextDetLocalVisualizer(**vis_cfg)
det_local_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')
det_local_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))
det_local_visualizer.add_datasample(
'image', image, self.data_sample, out_file=out_file)
self._assert_image_and_shape(out_file, (h, w * 2, c))
det_local_visualizer.add_datasample(
'image',
image,
self.data_sample,
draw_gt=False,
out_file=out_file)
self._assert_image_and_shape(out_file, (h, w, c))
det_local_visualizer.add_datasample(
'image',
image,
self.data_sample,
draw_pred=False,
out_file=out_file)
self._assert_image_and_shape(out_file, (h, w, c))
det_local_visualizer.add_datasample(
'image', image, None, out_file=out_file)
self._assert_image_and_shape(out_file, (h, w, 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)
|