From 5caa945a8dcd631d13efd7be7eb7c56fcc1ddcb3 Mon Sep 17 00:00:00 2001 From: Jianyong Chen <46100303+HolyCrap96@users.noreply.github.com> Date: Mon, 13 Dec 2021 17:45:30 +0800 Subject: [PATCH] fix #614: textsnake targets (#660) * fix #614: textsnake targets * fix lint * add textsnake_targets test cases * init with eps * fix test coverage --- .../textdet_targets/textsnake_targets.py | 100 ++++++++++-------- tests/test_dataset/test_textdet_targets.py | 21 ++++ 2 files changed, 76 insertions(+), 45 deletions(-) diff --git a/mmocr/datasets/pipelines/textdet_targets/textsnake_targets.py b/mmocr/datasets/pipelines/textdet_targets/textsnake_targets.py index c37d204f..3a8e4d21 100644 --- a/mmocr/datasets/pipelines/textdet_targets/textsnake_targets.py +++ b/mmocr/datasets/pipelines/textdet_targets/textsnake_targets.py @@ -32,30 +32,33 @@ class TextSnakeTargets(BaseTextDetTargets): self.orientation_thr = orientation_thr self.resample_step = resample_step self.center_region_shrink_ratio = center_region_shrink_ratio + self.eps = 1e-8 def vector_angle(self, vec1, vec2): if vec1.ndim > 1: - unit_vec1 = vec1 / (norm(vec1, axis=-1) + 1e-8).reshape((-1, 1)) + unit_vec1 = vec1 / (norm(vec1, axis=-1) + self.eps).reshape( + (-1, 1)) else: - unit_vec1 = vec1 / (norm(vec1, axis=-1) + 1e-8) + unit_vec1 = vec1 / (norm(vec1, axis=-1) + self.eps) if vec2.ndim > 1: - unit_vec2 = vec2 / (norm(vec2, axis=-1) + 1e-8).reshape((-1, 1)) + unit_vec2 = vec2 / (norm(vec2, axis=-1) + self.eps).reshape( + (-1, 1)) else: - unit_vec2 = vec2 / (norm(vec2, axis=-1) + 1e-8) + unit_vec2 = vec2 / (norm(vec2, axis=-1) + self.eps) return np.arccos( np.clip(np.sum(unit_vec1 * unit_vec2, axis=-1), -1.0, 1.0)) def vector_slope(self, vec): assert len(vec) == 2 - return abs(vec[1] / (vec[0] + 1e-8)) + return abs(vec[1] / (vec[0] + self.eps)) def vector_sin(self, vec): assert len(vec) == 2 - return vec[1] / (norm(vec) + 1e-8) + return vec[1] / (norm(vec) + self.eps) def vector_cos(self, vec): assert len(vec) == 2 - return vec[0] / (norm(vec) + 1e-8) + return vec[0] / (norm(vec) + self.eps) def find_head_tail(self, points, orientation_thr): """Find the head edge and tail edge of a text polygon. @@ -97,7 +100,7 @@ class TextSnakeTargets(BaseTextDetTargets): edge_dist = np.maximum( norm(pad_points[1:] - poly_center, axis=-1), norm(pad_points[:-1] - poly_center, axis=-1)) - dist_score = edge_dist / np.max(edge_dist) + dist_score = edge_dist / (np.max(edge_dist) + self.eps) position_score = np.zeros(len(edge_vec)) score = 0.5 * theta_sum_score + 0.15 * adjacent_theta_score score += 0.35 * dist_score @@ -198,6 +201,29 @@ class TextSnakeTargets(BaseTextDetTargets): return head_edge, tail_edge, top_sideline, bot_sideline + def cal_curve_length(self, line): + """Calculate the length of each edge on the discrete curve and the sum. + + Args: + line (ndarray): The points composing a discrete curve. + + Returns: + tuple: Returns (edges_length, total_length). + + - | edge_length (ndarray): The length of each edge on the + discrete curve. + - | total_length (float): The total length of the discrete + curve. + """ + + assert line.ndim == 2 + assert len(line) >= 2 + + edges_length = np.sqrt((line[1:, 0] - line[:-1, 0])**2 + + (line[1:, 1] - line[:-1, 1])**2) + total_length = np.sum(edges_length) + return edges_length, total_length + def resample_line(self, line, n): """Resample n points on a line. @@ -213,34 +239,24 @@ class TextSnakeTargets(BaseTextDetTargets): assert line.shape[0] >= 2 assert line.shape[1] == 2 assert isinstance(n, int) - assert n > 0 + assert n > 2 - length_list = [ - norm(line[i + 1] - line[i]) for i in range(len(line) - 1) - ] - total_length = sum(length_list) - length_cumsum = np.cumsum([0.0] + length_list) - delta_length = total_length / (float(n) + 1e-8) - - current_edge_ind = 0 - resampled_line = [line[0]] - - for i in range(1, n): - current_line_len = i * delta_length - - while current_line_len >= length_cumsum[current_edge_ind + 1]: - current_edge_ind += 1 - current_edge_end_shift = current_line_len - length_cumsum[ - current_edge_ind] - end_shift_ratio = current_edge_end_shift / length_list[ - current_edge_ind] - current_point = line[current_edge_ind] + ( - line[current_edge_ind + 1] - - line[current_edge_ind]) * end_shift_ratio - resampled_line.append(current_point) - - resampled_line.append(line[-1]) - resampled_line = np.array(resampled_line) + edges_length, total_length = self.cal_curve_length(line) + t_org = np.insert(np.cumsum(edges_length), 0, 0) + unit_t = total_length / (n - 1) + t_equidistant = np.arange(1, n - 1, dtype=np.float32) * unit_t + edge_ind = 0 + points = [line[0]] + for t in t_equidistant: + while edge_ind < len(edges_length) - 1 and t > t_org[edge_ind + 1]: + edge_ind += 1 + t_l, t_r = t_org[edge_ind], t_org[edge_ind + 1] + weight = np.array([t_r - t, t - t_l], dtype=np.float32) / ( + t_r - t_l + self.eps) + p_coords = np.dot(weight, line[[edge_ind, edge_ind + 1]]) + points.append(p_coords) + points.append(line[-1]) + resampled_line = np.vstack(points) return resampled_line @@ -266,17 +282,11 @@ class TextSnakeTargets(BaseTextDetTargets): assert sideline2.shape[0] >= 2 assert isinstance(resample_step, float) - length1 = sum([ - norm(sideline1[i + 1] - sideline1[i]) - for i in range(len(sideline1) - 1) - ]) - length2 = sum([ - norm(sideline2[i + 1] - sideline2[i]) - for i in range(len(sideline2) - 1) - ]) + _, length1 = self.cal_curve_length(sideline1) + _, length2 = self.cal_curve_length(sideline2) - total_length = (length1 + length2) / 2 - resample_point_num = max(int(float(total_length) / resample_step), 1) + avg_length = (length1 + length2) / 2 + resample_point_num = max(int(float(avg_length) / resample_step) + 1, 3) resampled_line1 = self.resample_line(sideline1, resample_point_num) resampled_line2 = self.resample_line(sideline2, resample_point_num) diff --git a/tests/test_dataset/test_textdet_targets.py b/tests/test_dataset/test_textdet_targets.py index f9c97bfe..2008c5c6 100644 --- a/tests/test_dataset/test_textdet_targets.py +++ b/tests/test_dataset/test_textdet_targets.py @@ -156,11 +156,21 @@ def test_gen_textsnake_targets(mock_show_feature): assert np.allclose(target_generator.resample_step, 4.0) assert np.allclose(target_generator.center_region_shrink_ratio, 0.3) + # test vector_angle + vec1 = np.array([[-1, 0], [0, 1]]) + vec2 = np.array([[1, 0], [0, 1]]) + angles = target_generator.vector_angle(vec1, vec2) + assert np.allclose(angles, np.array([np.pi, 0]), atol=1e-3) + # test find_head_tail for quadrangle polygon = np.array([[1.0, 1.0], [5.0, 1.0], [5.0, 3.0], [1.0, 3.0]]) head_inds, tail_inds = target_generator.find_head_tail(polygon, 2.0) assert np.allclose(head_inds, [3, 0]) assert np.allclose(tail_inds, [1, 2]) + polygon = np.array([[1.0, 1.0], [1.0, 3.0], [5.0, 3.0], [5.0, 1.0]]) + head_inds, tail_inds = target_generator.find_head_tail(polygon, 2.0) + assert np.allclose(head_inds, [0, 1]) + assert np.allclose(tail_inds, [2, 3]) # test find_head_tail for polygon polygon = np.array([[0., 10.], [3., 3.], [10., 0.], [17., 3.], [20., 10.], @@ -170,6 +180,17 @@ def test_gen_textsnake_targets(mock_show_feature): assert np.allclose(head_inds, [9, 0]) assert np.allclose(tail_inds, [4, 5]) + # test resample_line + line = np.array([[0, 0], [0, 1], [0, 3], [0, 4], [0, 7], [0, 8]]) + resampled_line = target_generator.resample_line(line, 3) + assert len(resampled_line) == 3 + assert np.allclose(resampled_line, np.array([[0, 0], [0, 4], [0, 8]])) + line = np.array([[0, 0], [0, 0]]) + resampled_line = target_generator.resample_line(line, 4) + assert len(resampled_line) == 4 + assert np.allclose(resampled_line, + np.array([[0, 0], [0, 0], [0, 0], [0, 0]])) + # test generate_text_region_mask img_size = (3, 10) text_polys = [[np.array([0, 0, 1, 0, 1, 1, 0, 1])],