mirror of https://github.com/open-mmlab/mmocr.git
267 lines
11 KiB
Python
267 lines
11 KiB
Python
# Copyright (c) OpenMMLab. All rights reserved.
|
|
import unittest
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
from mmocr.utils import (bbox2poly, bbox_center_distance, bbox_diag_distance,
|
|
bezier2polygon, is_on_same_line, rescale_bbox,
|
|
rescale_bboxes, stitch_boxes_into_lines)
|
|
from mmocr.utils.bbox_utils import bbox_jitter
|
|
|
|
|
|
class TestBbox2poly(unittest.TestCase):
|
|
|
|
def setUp(self) -> None:
|
|
self.box_array = np.array([1, 1, 2, 2])
|
|
self.box_list = [1, 1, 2, 2]
|
|
self.box_tensor = torch.tensor([1, 1, 2, 2])
|
|
self.gt_xyxy = np.array([1, 1, 2, 1, 2, 2, 1, 2])
|
|
self.gt_xywh = np.array([1, 1, 3, 1, 3, 3, 1, 3])
|
|
|
|
def test_bbox2poly(self):
|
|
# mode: xyxy
|
|
# test np.array
|
|
self.assertTrue(
|
|
np.array_equal(bbox2poly(self.box_array), self.gt_xyxy))
|
|
# test list
|
|
self.assertTrue(np.array_equal(bbox2poly(self.box_list), self.gt_xyxy))
|
|
# test tensor
|
|
self.assertTrue(
|
|
np.array_equal(bbox2poly(self.box_tensor), self.gt_xyxy))
|
|
|
|
# mode: xywh
|
|
# test np.array
|
|
self.assertTrue(
|
|
np.array_equal(
|
|
bbox2poly(self.box_array, mode='xywh'), self.gt_xywh))
|
|
# test list
|
|
self.assertTrue(
|
|
np.array_equal(
|
|
bbox2poly(self.box_list, mode='xywh'), self.gt_xywh))
|
|
# test tensor
|
|
self.assertTrue(
|
|
np.array_equal(
|
|
bbox2poly(self.box_tensor, mode='xywh'), self.gt_xywh))
|
|
|
|
# invalid mode
|
|
with self.assertRaises(NotImplementedError):
|
|
bbox2poly(self.box_tensor, mode='a')
|
|
|
|
|
|
class TestBoxCenterDistance(unittest.TestCase):
|
|
|
|
def setUp(self) -> None:
|
|
self.box1_list = [1, 1, 3, 3]
|
|
self.box2_list = [2, 2, 4, 2]
|
|
self.box1_array = np.array([1, 1, 3, 3])
|
|
self.box2_array = np.array([2, 2, 4, 2])
|
|
self.box1_tensor = torch.tensor([1, 1, 3, 3])
|
|
self.box2_tensor = torch.tensor([2, 2, 4, 2])
|
|
self.gt = 1
|
|
|
|
def test_box_center_distance(self):
|
|
# test list
|
|
self.assertEqual(
|
|
bbox_center_distance(self.box1_list, self.box2_list), self.gt)
|
|
# test np.array
|
|
self.assertEqual(
|
|
bbox_center_distance(self.box1_array, self.box2_array), self.gt)
|
|
# test tensor
|
|
self.assertEqual(
|
|
bbox_center_distance(self.box1_tensor, self.box2_tensor), self.gt)
|
|
|
|
|
|
class TestBoxDiagDistance(unittest.TestCase):
|
|
|
|
def setUp(self) -> None:
|
|
self.box_list1 = [0, 0, 1, 1, 0, 10, -10, 0]
|
|
self.box_array1 = np.array(self.box_list1)
|
|
self.box_tensor1 = torch.tensor(self.box_list1)
|
|
self.gt1 = 10
|
|
self.box_list2 = [0, 0, 1, 1]
|
|
self.box_array2 = np.array(self.box_list2)
|
|
self.box_tensor2 = torch.tensor(self.box_list2)
|
|
self.gt2 = np.sqrt(2)
|
|
|
|
def test_bbox_diag_distance(self):
|
|
# quad [x1, y1, x2, y2, x3, y3, x4, y4]
|
|
# list
|
|
self.assertEqual(bbox_diag_distance(self.box_list1), self.gt1)
|
|
# array
|
|
self.assertEqual(bbox_diag_distance(self.box_array1), self.gt1)
|
|
# tensor
|
|
self.assertEqual(bbox_diag_distance(self.box_tensor1), self.gt1)
|
|
# rect [x1, y1, x2, y2]
|
|
# list
|
|
self.assertAlmostEqual(bbox_diag_distance(self.box_list2), self.gt2)
|
|
# array
|
|
self.assertAlmostEqual(bbox_diag_distance(self.box_array2), self.gt2)
|
|
# tensor
|
|
self.assertAlmostEqual(bbox_diag_distance(self.box_tensor2), self.gt2)
|
|
|
|
|
|
class TestBezier2Polygon(unittest.TestCase):
|
|
|
|
def setUp(self) -> None:
|
|
self.bezier_points1 = [
|
|
37.0, 249.0, 72.5, 229.55, 95.34, 220.65, 134.0, 216.0, 132.0,
|
|
233.0, 82.11, 240.2, 72.46, 247.16, 38.0, 263.0
|
|
]
|
|
self.gt1 = np.array([[37.0, 249.0],
|
|
[42.50420761043885, 246.01570199737577],
|
|
[47.82291296107305, 243.2012392477038],
|
|
[52.98102930456334, 240.5511007435486],
|
|
[58.00346989357049, 238.05977547747486],
|
|
[62.91514798075522, 235.721752442047],
|
|
[67.74097681877824, 233.53152062982943],
|
|
[72.50586966030032, 231.48356903338674],
|
|
[77.23473975798221, 229.57238664528356],
|
|
[81.95250036448464, 227.79246245808432],
|
|
[86.68406473246829, 226.13828546435346],
|
|
[91.45434611459396, 224.60434465665548],
|
|
[96.28825776352238, 223.18512902755504],
|
|
[101.21071293191426, 221.87512756961655],
|
|
[106.24662487243039, 220.6688292754046],
|
|
[111.42090683773145, 219.5607231374836],
|
|
[116.75847208047819, 218.5452981484181],
|
|
[122.28423385333137, 217.6170433007727],
|
|
[128.02310540895172, 216.77044758711182],
|
|
[134.0, 216.0], [132.0, 233.0],
|
|
[124.4475521213005, 234.13617728531858],
|
|
[117.50700976818779, 235.2763434903047],
|
|
[111.12146960198277, 236.42847645429362],
|
|
[105.2340282840064, 237.6005540166205],
|
|
[99.78778247557953, 238.80055401662054],
|
|
[94.72582883802303, 240.0364542936288],
|
|
[89.99126403265781, 241.31623268698053],
|
|
[85.52718472080478, 242.64786703601104],
|
|
[81.27668756378483, 244.03933518005545],
|
|
[77.1828692229188, 245.49861495844874],
|
|
[73.18882635952762, 247.0336842105263],
|
|
[69.23765563493221, 248.65252077562326],
|
|
[65.27245371045342, 250.3631024930748],
|
|
[61.23631724741216, 252.17340720221605],
|
|
[57.07234290712931, 254.09141274238226],
|
|
[52.723627350925796, 256.12509695290856],
|
|
[48.13326724012247, 258.2824376731302],
|
|
[43.24435923604024, 260.5714127423822],
|
|
[38.0, 263.0]])
|
|
self.bezier_points2 = [0, 0, 0, 1, 0, 2, 0, 3, 1, 0, 1, 1, 1, 2, 1, 3]
|
|
self.gt2 = np.array([[0, 0], [0, 1.5], [0, 3], [1, 0], [1, 1.5],
|
|
[1, 3]])
|
|
self.invalid_input = [0, 1]
|
|
|
|
def test_bezier2polygon(self):
|
|
self.assertTrue(
|
|
np.allclose(bezier2polygon(self.bezier_points1), self.gt1))
|
|
with self.assertRaises(AssertionError):
|
|
bezier2polygon(self.bezier_points2, num_sample=-1)
|
|
with self.assertRaises(AssertionError):
|
|
bezier2polygon(self.invalid_input, num_sample=-1)
|
|
|
|
|
|
class TestBboxJitter(unittest.TestCase):
|
|
|
|
def test_bbox_jitter(self):
|
|
dummy_points_x = [20, 120, 120, 20]
|
|
dummy_points_y = [20, 20, 40, 40]
|
|
|
|
kwargs = dict(jitter_ratio_x=0.0, jitter_ratio_y=0.0)
|
|
|
|
with self.assertRaises(AssertionError):
|
|
bbox_jitter([], dummy_points_y)
|
|
with self.assertRaises(AssertionError):
|
|
bbox_jitter(dummy_points_x, [])
|
|
with self.assertRaises(AssertionError):
|
|
bbox_jitter(dummy_points_x, dummy_points_y, jitter_ratio_x=1.)
|
|
with self.assertRaises(AssertionError):
|
|
bbox_jitter(dummy_points_x, dummy_points_y, jitter_ratio_y=1.)
|
|
|
|
bbox_jitter(dummy_points_x, dummy_points_y, **kwargs)
|
|
|
|
assert np.allclose(dummy_points_x, [20, 120, 120, 20])
|
|
assert np.allclose(dummy_points_y, [20, 20, 40, 40])
|
|
|
|
|
|
class TestIsOnSameLine(unittest.TestCase):
|
|
|
|
def test_box_on_line(self):
|
|
# regular boxes
|
|
box1 = [0, 0, 1, 0, 1, 1, 0, 1]
|
|
box2 = [2, 0.5, 3, 0.5, 3, 1.5, 2, 1.5]
|
|
box3 = [4, 0.8, 5, 0.8, 5, 1.8, 4, 1.8]
|
|
self.assertTrue(is_on_same_line(box1, box2, 0.5))
|
|
self.assertFalse(is_on_same_line(box1, box3, 0.5))
|
|
|
|
# irregular box4
|
|
box4 = [0, 0, 1, 1, 1, 2, 0, 1]
|
|
box5 = [2, 1.5, 3, 1.5, 3, 2.5, 2, 2.5]
|
|
box6 = [2, 1.6, 3, 1.6, 3, 2.6, 2, 2.6]
|
|
self.assertTrue(is_on_same_line(box4, box5, 0.5))
|
|
self.assertFalse(is_on_same_line(box4, box6, 0.5))
|
|
|
|
|
|
class TestStitchBoxesIntoLines(unittest.TestCase):
|
|
|
|
def test_stitch_boxes_into_lines(self):
|
|
boxes = [ # regular boxes
|
|
[0, 0, 1, 0, 1, 1, 0, 1],
|
|
[2, 0.5, 3, 0.5, 3, 1.5, 2, 1.5],
|
|
[3, 1.2, 4, 1.2, 4, 2.2, 3, 2.2],
|
|
[5, 0.5, 6, 0.5, 6, 1.5, 5, 1.5],
|
|
# irregular box
|
|
[6, 1.5, 7, 1.25, 7, 1.75, 6, 1.75]
|
|
]
|
|
raw_input = [{
|
|
'box': boxes[i],
|
|
'text': str(i)
|
|
} for i in range(len(boxes))]
|
|
result = stitch_boxes_into_lines(raw_input, 1, 0.5)
|
|
# Final lines: [0, 1], [2], [3, 4]
|
|
# box 0, 1, 3, 4 are on the same line but box 3 is 2 pixels away from
|
|
# box 1
|
|
# box 3 and 4 are on the same line since the length of overlapping part
|
|
# >= 0.5 * the y-axis length of box 5
|
|
expected_result = [{
|
|
'box': [0, 0, 3, 0, 3, 1.5, 0, 1.5],
|
|
'text': '0 1'
|
|
}, {
|
|
'box': [3, 1.2, 4, 1.2, 4, 2.2, 3, 2.2],
|
|
'text': '2'
|
|
}, {
|
|
'box': [5, 0.5, 7, 0.5, 7, 1.75, 5, 1.75],
|
|
'text': '3 4'
|
|
}]
|
|
result.sort(key=lambda x: x['box'][0])
|
|
expected_result.sort(key=lambda x: x['box'][0])
|
|
self.assertEqual(result, expected_result)
|
|
|
|
|
|
class TestRescaleBbox(unittest.TestCase):
|
|
|
|
def setUp(self) -> None:
|
|
self.bbox = np.array([0, 0, 1, 1])
|
|
self.bboxes = np.array([[0, 0, 1, 1], [1, 1, 2, 2]])
|
|
self.scale = 2
|
|
|
|
def test_rescale_bbox(self):
|
|
# mul
|
|
rescaled_bbox = rescale_bbox(self.bbox, self.scale, mode='mul')
|
|
self.assertTrue(np.allclose(rescaled_bbox, np.array([0, 0, 2, 2])))
|
|
# div
|
|
rescaled_bbox = rescale_bbox(self.bbox, self.scale, mode='div')
|
|
self.assertTrue(np.allclose(rescaled_bbox, np.array([0, 0, 0.5, 0.5])))
|
|
|
|
def test_rescale_bboxes(self):
|
|
# mul
|
|
rescaled_bboxes = rescale_bboxes(self.bboxes, self.scale, mode='mul')
|
|
self.assertTrue(
|
|
np.allclose(rescaled_bboxes, np.array([[0, 0, 2, 2], [2, 2, 4,
|
|
4]])))
|
|
# div
|
|
rescaled_bboxes = rescale_bboxes(self.bboxes, self.scale, mode='div')
|
|
self.assertTrue(
|
|
np.allclose(rescaled_bboxes,
|
|
np.array([[0, 0, 0.5, 0.5], [0.5, 0.5, 1, 1]])))
|