def generate_rbox(im_size, polys, tags):
h, w = im_size
poly_mask = np.zeros((h, w), dtype=np.uint8)
score_map = np.zeros((h, w), dtype=np.uint8)
geo_map = np.zeros((h, w, 8), dtype=np.float32)
# mask used during traning, to ignore some hard areas
training_mask = np.ones((h, w), dtype=np.uint8)
for poly_idx, poly_tag in enumerate(zip(polys, tags)):
poly = poly_tag[0]
tag = poly_tag[1]
r = [None, None, None, None]
for i in range(4):
r[i] = min(np.linalg.norm(poly[i] - poly[(i + 1) % 4]),
np.linalg.norm(poly[i] - poly[(i - 1) % 4]))
# score map
# shrinked_poly = shrink_poly(poly.copy(), r).astype(np.int32)[np.newaxis, :, :]
# close shrink function
shrinked_poly = poly.astype(np.int32)[np.newaxis, :,:]
cv2.fillPoly(score_map, shrinked_poly, 1)
cv2.fillPoly(poly_mask, shrinked_poly, poly_idx + 1)
# if the poly is too small, then ignore it during training
poly_h = min(np.linalg.norm(poly[0] - poly[3]), np.linalg.norm(poly[1] - poly[2]))
poly_w = min(np.linalg.norm(poly[0] - poly[1]), np.linalg.norm(poly[2] - poly[3]))
if min(poly_h, poly_w) < FLAGS.min_text_size:
cv2.fillPoly(training_mask, poly.astype(np.int32)[np.newaxis, :, :], 0)
if tag:
cv2.fillPoly(training_mask, poly.astype(np.int32)[np.newaxis, :, :], 0)
xy_in_poly = np.argwhere(poly_mask == (poly_idx + 1))
for y, x in xy_in_poly:
point = np.array([x, y], dtype=np.int32)
# left
geo_map[y, x, 0] = valid_link(point, score_map, w, h,'left')
# left_down
geo_map[y, x, 1] = valid_link(point, score_map, w, h, 'left_down')
# left_up
geo_map[y, x, 2] = valid_link(point, score_map, w, h, 'left_up')
# right
geo_map[y, x, 3] = valid_link(point, score_map, w, h,'right')
# right_down
geo_map[y, x, 4] = valid_link(point, score_map, w, h,'right_down')
# right_up
geo_map[y, x, 5] = valid_link(point, score_map, w, h, 'right_up')
# up
geo_map[y, x, 6] = valid_link(point, score_map, w, h, 'up')
# down
geo_map[y, x, 7] = valid_link(point, score_map, w, h, 'down')
return score_map, geo_map, training_mask
评论列表
文章目录
