Mkae sort_vertex public api

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

mmocr/datasets/kie_dataset.py

@@ -171,7 +171,7 @@ class KIEDataset(BaseDataset):
         }
 
     def list_to_numpy(self, ann_infos):
-        """Convert list to np.ndarray."""
+        """Convert bboxes, relations, texts and labels to ndarray."""
         boxes, text_inds = ann_infos['boxes'], ann_infos['text_inds']
         boxes = np.array(boxes, np.int32)
         relations, bboxes = self.compute_relation(boxes)
@@ -188,7 +188,7 @@ class KIEDataset(BaseDataset):
                     edges = (edges & labels == 1).astype(np.int32)
                 np.fill_diagonal(edges, -1)
                 labels = np.concatenate([labels, edges], -1)
-        padded_text_inds = self.pad_text_ind(text_inds)
+        padded_text_inds = self.pad_text_indices(text_inds)
 
         return dict(
             bboxes=bboxes,
@@ -196,7 +196,7 @@ class KIEDataset(BaseDataset):
             texts=padded_text_inds,
             labels=labels)
 
-    def pad_text_ind(self, text_inds):
+    def pad_text_indices(self, text_inds):
         """Pad text index to same length."""
         max_len = max([len(text_ind) for text_ind in text_inds])
         padded_text_inds = -np.ones((len(text_inds), max_len), np.int32)

mmocr/datasets/pipelines/__init__.py

@@ -1,3 +1,4 @@
+from .box_utils import sort_vertex
 from .custom_format_bundle import CustomFormatBundle
 from .dbnet_transforms import EastRandomCrop, ImgAug
 from .kie_transforms import KIEFormatBundle
@@ -22,5 +23,5 @@ __all__ = [
     'DRRGTargets', 'RandomCropPolyInstances', 'RandomRotatePolyInstances',
     'RandomPaddingOCR', 'ImgAug', 'EastRandomCrop', 'RandomRotateImageBox',
     'OpencvToPil', 'PilToOpencv', 'KIEFormatBundle', 'SquareResizePad',
-    'TextSnakeTargets'
+    'TextSnakeTargets', 'sort_vertex'
 ]

mmocr/datasets/pipelines/box_utils.py

@@ -0,0 +1,81 @@
+import mmocr.utils as utils
+import numpy as np
+from shapely.geometry import LineString, Point, Polygon
+
+def sort_vertex(points_x, points_y):
+    """Sort box vertices in clockwise order from left-top 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
+
+    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)
+    y_arr = np.array(y - center_y)
+
+    angle = np.arctan2(y_arr, x_arr) * 180.0 / np.pi
+    sort_idx = np.argsort(angle)
+
+    sorted_points_x, sorted_points_y = [], []
+    for i in range(4):
+        sorted_points_x.append(points_x[sort_idx[i]])
+        sorted_points_y.append(points_y[sort_idx[i]])
+
+    return convert_canonical(sorted_points_x, sorted_points_y)
+
+
+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

mmocr/datasets/pipelines/crop.py

@@ -1,87 +1,9 @@
 import cv2
 import numpy as np
-from shapely.geometry import LineString, Point, Polygon
+from shapely.geometry import LineString, Point
 
 import mmocr.utils as utils
-
-
-def sort_vertex(points_x, points_y):
-    """Sort box vertices in clockwise order from left-top 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
-
-    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)
-    y_arr = np.array(y - center_y)
-
-    angle = np.arctan2(y_arr, x_arr) * 180.0 / np.pi
-    sort_idx = np.argsort(angle)
-
-    sorted_points_x, sorted_points_y = [], []
-    for i in range(4):
-        sorted_points_x.append(points_x[sort_idx[i]])
-        sorted_points_y.append(points_y[sort_idx[i]])
-
-    return convert_canonical(sorted_points_x, sorted_points_y)
-
-
-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
+from .box_utils import sort_vertex
 
 
 def box_jitter(points_x, points_y, jitter_ratio_x=0.5, jitter_ratio_y=0.1):