Speed up sort_vertex (#239)

* Speed up sort_vertex

Signed-off-by: lizz <lizz@sensetime.com>

* Fix test

Signed-off-by: lizz <lizz@sensetime.com>

mmocr/datasets/pipelines/box_utils.py

@@ -1,5 +1,4 @@
 import numpy as np
-from shapely.geometry import LineString, Point, Polygon
 
 import mmocr.utils as utils
 
@@ -20,76 +19,36 @@ def sort_vertex(points_x, points_y):
         points_y, int)
     assert len(points_x) == 4
     assert len(points_y) == 4
+    vertices = np.stack((points_x, points_y), axis=-1).astype(np.float32)
+    vertices = _sort_vertex(vertices)
+    sorted_points_x = list(vertices[:, 0])
+    sorted_points_y = list(vertices[:, 1])
+    return sorted_points_x, sorted_points_y
 
-    x = np.array(points_x)
-    y = np.array(points_y)
-    center_x = np.sum(x) * 0.25
-    center_y = np.sum(y) * 0.25
 
-    x_arr = np.array(x - center_x)
+def _sort_vertex(vertices):
-    y_arr = np.array(y - center_y)
+    assert vertices.ndim == 2
+    assert vertices.shape[-1] == 2
+    N = vertices.shape[0]
+    if N == 0:
+        return vertices
 
-    angle = np.arctan2(y_arr, x_arr) * 180.0 / np.pi
+    center = np.mean(vertices, axis=0)
-    sort_idx = np.argsort(angle)
+    directions = vertices - center
+    angles = np.arctan2(directions[:, 1], directions[:, 0])
+    sort_idx = np.argsort(angles)
+    vertices = vertices[sort_idx]
 
-    sorted_points_x, sorted_points_y = [], []
+    left_top = np.min(vertices, axis=0)
-    for i in range(4):
+    dists = np.linalg.norm(left_top - vertices, axis=-1, ord=2)
-        sorted_points_x.append(points_x[sort_idx[i]])
+    lefttop_idx = np.argmin(dists)
-        sorted_points_y.append(points_y[sort_idx[i]])
+    indexes = (np.arange(N, dtype=np.int) + lefttop_idx) % N
-
+    return vertices[indexes]
-    return convert_canonical(sorted_points_x, sorted_points_y)
 
 
 def sort_vertex8(points):
     """Sort vertex with 8 points [x1 y1 x2 y2 x3 y3 x4 y4]"""
     assert len(points) == 8
-    x_list, y_list = points[0::2], points[1::2]
+    vertices = _sort_vertex(np.array(points, dtype=np.float32).reshape(-1, 2))
-    sorted_x_list, sorted_y_list = sort_vertex(x_list, y_list)
+    sorted_box = list(vertices.flatten())
-    sorted_box = []
-    for x, y in zip(sorted_x_list, sorted_y_list):
-        sorted_box.append(x)
-        sorted_box.append(y)
     return sorted_box
-
-
-def convert_canonical(points_x, points_y):
-    """Make left-top be first.
-
-    Args:
-        points_x (list[float]): x of four vertices.
-        points_y (list[float]): y of four vertices.
-    Returns:
-        sorted_points_x (list[float]): x of sorted four vertices.
-        sorted_points_y (list[float]): y of sorted four vertices.
-    """
-    assert utils.is_type_list(points_x, float) or utils.is_type_list(
-        points_x, int)
-    assert utils.is_type_list(points_y, float) or utils.is_type_list(
-        points_y, int)
-    assert len(points_x) == 4
-    assert len(points_y) == 4
-
-    points = [Point(points_x[i], points_y[i]) for i in range(4)]
-
-    polygon = Polygon([(p.x, p.y) for p in points])
-    min_x, min_y, _, _ = polygon.bounds
-    points_to_lefttop = [
-        LineString([points[i], Point(min_x, min_y)]) for i in range(4)
-    ]
-    distances = np.array([line.length for line in points_to_lefttop])
-    sort_dist_idx = np.argsort(distances)
-    lefttop_idx = sort_dist_idx[0]
-
-    if lefttop_idx == 0:
-        point_orders = [0, 1, 2, 3]
-    elif lefttop_idx == 1:
-        point_orders = [1, 2, 3, 0]
-    elif lefttop_idx == 2:
-        point_orders = [2, 3, 0, 1]
-    else:
-        point_orders = [3, 0, 1, 2]
-
-    sorted_points_x = [points_x[i] for i in point_orders]
-    sorted_points_y = [points_y[j] for j in point_orders]
-
-    return sorted_points_x, sorted_points_y

tests/test_dataset/test_crop.py

@@ -1,9 +1,10 @@
 import math
+from itertools import chain, permutations
 
 import numpy as np
 import pytest
 
-from mmocr.datasets.pipelines.box_utils import convert_canonical, sort_vertex
+from mmocr.datasets.pipelines.box_utils import sort_vertex, sort_vertex8
 from mmocr.datasets.pipelines.crop import box_jitter, crop_img, warp_img
 
 
@@ -11,38 +12,36 @@ def test_order_vertex():
     dummy_points_x = [20, 20, 120, 120]
     dummy_points_y = [20, 40, 40, 20]
 
+    expect_points_x = [20, 120, 120, 20]
+    expect_points_y = [20, 20, 40, 40]
+
     with pytest.raises(AssertionError):
         sort_vertex([], dummy_points_y)
     with pytest.raises(AssertionError):
         sort_vertex(dummy_points_x, [])
 
-    ordered_points_x, ordered_points_y = sort_vertex(dummy_points_x,
+    for perm in set(permutations([0, 1, 2, 3])):
-                                                     dummy_points_y)
+        points_x = [dummy_points_x[i] for i in perm]
-
+        points_y = [dummy_points_y[i] for i in perm]
-    expect_points_x = [20, 120, 120, 20]
+        ordered_points_x, ordered_points_y = sort_vertex(points_x, points_y)
-    expect_points_y = [20, 20, 40, 40]
 
         assert np.allclose(ordered_points_x, expect_points_x)
         assert np.allclose(ordered_points_y, expect_points_y)
 
 
-def test_convert_canonical():
+def test_sort_vertex8():
-    dummy_points_x = [120, 120, 20, 20]
+    dummy_points_x = [21, 21, 122, 122]
-    dummy_points_y = [20, 40, 40, 20]
+    dummy_points_y = [21, 39, 39, 21]
 
-    with pytest.raises(AssertionError):
+    expect_points = [21, 21, 122, 21, 122, 39, 21, 39]
-        convert_canonical([], dummy_points_y)
-    with pytest.raises(AssertionError):
-        convert_canonical(dummy_points_x, [])
 
-    ordered_points_x, ordered_points_y = convert_canonical(
+    for perm in set(permutations([0, 1, 2, 3])):
-        dummy_points_x, dummy_points_y)
+        points_x = [dummy_points_x[i] for i in perm]
+        points_y = [dummy_points_y[i] for i in perm]
+        points = list(chain.from_iterable(zip(points_x, points_y)))
+        ordered_points = sort_vertex8(points)
 
-    expect_points_x = [20, 120, 120, 20]
+        assert np.allclose(ordered_points, expect_points)
-    expect_points_y = [20, 20, 40, 40]
-
-    assert np.allclose(ordered_points_x, expect_points_x)
-    assert np.allclose(ordered_points_y, expect_points_y)
 
 
 def test_box_jitter():