mirrors
/
mmocr
mirror of https://github.com/open-mmlab/mmocr.git
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

fix #614: textsnake targets (#660) 

* fix #614: textsnake targets

* fix lint

* add textsnake_targets test cases

* init with eps

* fix test coverage
pull/651/head
Jianyong Chen 2021-12-13 17:45:30 +08:00 committed by GitHub
parent b04775fd78
commit 5caa945a8d
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 76 additions and 45 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

100
mmocr/datasets/pipelines/textdet_targets/textsnake_targets.py
View File

@ -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)

21
tests/test_dataset/test_textdet_targets.py
View File

@ -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])],