[Feature] Add bezier_to_polygon and sort_points to box_util (#703)

mmocr/utils/__init__.py

@@ -1,7 +1,8 @@
 # Copyright (c) OpenMMLab. All rights reserved.
 from mmcv.utils import Registry, build_from_cfg
 
-from .box_util import is_on_same_line, stitch_boxes_into_lines
+from .box_util import (bezier_to_polygon, is_on_same_line, sort_points,
+                       stitch_boxes_into_lines)
 from .check_argument import (equal_len, is_2dlist, is_3dlist, is_none_or_type,
                              is_type_list, valid_boundary)
 from .collect_env import collect_env
@@ -19,5 +20,5 @@
     'valid_boundary', 'lmdb_converter', 'drop_orientation',
     'convert_annotations', 'is_not_png', 'list_to_file', 'list_from_file',
     'is_on_same_line', 'stitch_boxes_into_lines', 'StringStrip',
-    'revert_sync_batchnorm'
+    'revert_sync_batchnorm', 'bezier_to_polygon', 'sort_points'
 ]

mmocr/utils/box_util.py

@@ -1,6 +1,10 @@
 # Copyright (c) OpenMMLab. All rights reserved.
+import functools
+
 import numpy as np
 
+from mmocr.utils.check_argument import is_2dlist, is_type_list
+
 
 def is_on_same_line(box_a, box_b, min_y_overlap_ratio=0.8):
     """Check if two boxes are on the same line by their y-axis coordinates.
@@ -116,3 +120,80 @@ def stitch_boxes_into_lines(boxes, max_x_dist=10, min_y_overlap_ratio=0.8):
             merged_boxes.append(merged_box)
 
     return merged_boxes
+
+
+def bezier_to_polygon(bezier_points, num_sample=20):
+    """Sample points from the boundary of a polygon enclosed by two Bezier
+    curves, which are controlled by ``bezier_points``.
+
+    Args:
+        bezier_points (ndarray): A :math:`(2, 4, 2)` array of 8 Bezeir points
+            or its equalivance. The first 4 points control the curve at one
+            side and the last four control the other side.
+        num_sample (int): The number of sample points at each Bezeir curve.
+
+    Returns:
+        list[ndarray]: A list of 2*num_sample points representing the polygon
+        extracted from Bezier curves.
+
+    Warning:
+        The points are not guaranteed to be ordered. Please use
+        :func:`mmocr.utils.sort_points` to sort points if necessary.
+    """
+    assert num_sample > 0
+
+    bezier_points = np.asarray(bezier_points)
+    assert np.prod(
+        bezier_points.shape) == 16, 'Need 8 Bezier control points to continue!'
+
+    bezier = bezier_points.reshape(2, 4, 2).transpose(0, 2, 1).reshape(4, 4)
+    u = np.linspace(0, 1, num_sample)
+
+    points = np.outer((1 - u) ** 3, bezier[:, 0]) \
+        + np.outer(3 * u * ((1 - u) ** 2), bezier[:, 1]) \
+        + np.outer(3 * (u ** 2) * (1 - u), bezier[:, 2]) \
+        + np.outer(u ** 3, bezier[:, 3])
+
+    # Convert points to polygon
+    points = np.concatenate((points[:, :2], points[:, 2:]), axis=0)
+    return points.tolist()
+
+
+def sort_points(points):
+    """Sort arbitory points in clockwise order. Reference:
+    https://stackoverflow.com/a/6989383.
+
+    Args:
+        points (list[ndarray] or ndarray or list[list]): A list of unsorted
+            boundary points.
+
+    Returns:
+        list[ndarray]: A list of points sorted in clockwise order.
+    """
+
+    assert is_type_list(points, np.ndarray) or isinstance(points, np.ndarray) \
+        or is_2dlist(points)
+
+    points = np.array(points)
+    center = np.mean(points, axis=0)
+
+    def cmp(a, b):
+        oa = a - center
+        ob = b - center
+
+        # Some corner cases
+        if oa[0] >= 0 and ob[0] < 0:
+            return 1
+        if oa[0] < 0 and ob[0] >= 0:
+            return -1
+
+        prod = np.cross(oa, ob)
+        if prod > 0:
+            return 1
+        if prod < 0:
+            return -1
+
+        # a, b are on the same line from the center
+        return 1 if (oa**2).sum() < (ob**2).sum() else -1
+
+    return sorted(points, key=functools.cmp_to_key(cmp))

tests/test_utils/test_box.py

@@ -1,5 +1,9 @@
 # Copyright (c) OpenMMLab. All rights reserved.
-from mmocr.utils import is_on_same_line, stitch_boxes_into_lines
+import numpy as np
+import pytest
+
+from mmocr.utils import (bezier_to_polygon, is_on_same_line, sort_points,
+                         stitch_boxes_into_lines)
 
 
 def test_box_on_line():
@@ -46,3 +50,79 @@ def test_stitch_boxes_into_lines():
     result.sort(key=lambda x: x['box'][0])
     expected_result.sort(key=lambda x: x['box'][0])
     assert result == expected_result
+
+
+def test_bezier_to_polygon():
+    bezier_points = [
+        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
+    ]
+    pts = bezier_to_polygon(bezier_points)
+    target = 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]])
+    assert np.allclose(pts, target)
+
+    bezier_points = [0, 0, 0, 1, 0, 2, 0, 3, 1, 0, 1, 1, 1, 2, 1, 3]
+    pts = bezier_to_polygon(bezier_points, num_sample=3)
+    target = np.array([[0, 0], [0, 1.5], [0, 3], [1, 0], [1, 1.5], [1, 3]])
+    assert np.allclose(pts, target)
+
+    with pytest.raises(AssertionError):
+        bezier_to_polygon(bezier_points, num_sample=-1)
+
+    bezier_points = [0, 1]
+    with pytest.raises(AssertionError):
+        bezier_to_polygon(bezier_points)
+
+
+def test_sort_points():
+    points = np.array([[1, 1], [0, 0], [1, -1], [2, -2], [0, 2], [1, 1],
+                       [0, 1], [-1, 1], [-1, -1]])
+    target = np.array([[-1, -1], [0, 0], [-1, 1], [0, 1], [0, 2], [1, 1],
+                       [1, 1], [2, -2], [1, -1]])
+    assert np.allclose(target, sort_points(points))
+
+    points = np.array([[1, 1], [1, -1], [-1, 1], [-1, -1]])
+    target = np.array([[-1, -1], [-1, 1], [1, 1], [1, -1]])
+    assert np.allclose(target, sort_points(points))
+
+    points = [[1, 1], [1, -1], [-1, 1], [-1, -1]]
+    assert np.allclose(target, sort_points(points))
+
+    with pytest.raises(AssertionError):
+        sort_points([1, 2])