def build_tcl(self, tcl_pred, sin_pred, cos_pred, radii_pred):
"""
Find TCL's center points and radii of each point
:param tcl_pred: output tcl mask, (512, 512)
:param sin_pred: output sin map, (512, 512)
:param cos_pred: output cos map, (512, 512)
:param radii_pred: output radii map, (512, 512)
:return: (list), tcl array: (n, 3), 3 denotes (x, y, radii)
"""
all_tcls = []
# find disjoint regions
tcl_mask = fill_hole(tcl_pred)
_, tcl_contours, _ = cv2.findContours(tcl_mask.astype(np.uint8), cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
for cont in tcl_contours:
# find an inner point of polygon
init = self.find_innerpoint(cont)
if init is None:
continue
x_init, y_init = init
# find left/right tcl
tcl_left = self.mask_to_tcl(sin_pred, cos_pred, radii_pred, cont, (x_init, y_init), direct=1)
tcl_right = self.mask_to_tcl(sin_pred, cos_pred, radii_pred, cont, (x_init, y_init), direct=-1)
# concat
tcl = np.concatenate([tcl_left[::-1][:-1], tcl_right])
all_tcls.append(tcl)
return all_tcls
def proposals_layer(self, tr_map, tcl_map, radii_map, sin_map, cos_map):
tr_pred_mask = tr_map > self.tr_thresh
tcl_pred_mask = tcl_map > self.tcl_thresh
# multiply TR and TCL
tcl_mask = tcl_pred_mask * tr_pred_mask
# regularize
sin_map, cos_map = regularize_sin_cos(sin_map, cos_map)
# find disjoint regions
tcl_mask = fill_hole(tcl_mask)
tcl_contours, _ = cv2.findContours(tcl_mask.astype(np.uint8),
cv2.RETR_TREE,
cv2.CHAIN_APPROX_SIMPLE)
mask = np.zeros_like(tcl_mask)
proposals = None
for cont in tcl_contours:
deal_map = mask.copy()
cv2.drawContours(deal_map, [cont], -1, 1, -1)
if deal_map.sum() <= 100:
continue
text_map = tr_map * deal_map
# ## 1. Reverse generation of box
bboxs = bbox_transfor_inv(radii_map,
sin_map,
cos_map,
text_map,
wclip=self.clip,
expend=self.expend)
# ## 3. local nms
bboxs = lanms.merge_quadrangle_n9(bboxs.astype('float32'), 0.25)
reconstruct_mask = mask.copy()
boxes = bboxs[:, :8].reshape((-1, 4, 2)).astype(np.int32)
cv2.drawContours(reconstruct_mask, boxes, -1, 1, -1)
if (reconstruct_mask *
tr_pred_mask).sum() < reconstruct_mask.sum() * 0.5:
continue
if proposals is None:
proposals = bboxs
else:
proposals = np.concatenate([proposals, bboxs], axis=0)
if proposals is None or proposals.shape[0] <= 0:
return None, None
# ## 5. generate cluster label
cxy = np.mean(proposals[:, :8].reshape((-1, 4, 2)),
axis=1).astype(np.int32)
# ## 6. Geometric features
gh = (radii_map[:, :, 0] + radii_map[:, :, 1])
h_map = gh[cxy[:, 1], cxy[:, 0]]
w_map = np.clip(h_map // 2, 2 * self.clip[0], 2 * self.clip[1])
c_map = cos_map[cxy[:, 1], cxy[:, 0]]
s_map = sin_map[cxy[:, 1], cxy[:, 0]]
geo_map = np.stack([cxy[:, 0], cxy[:, 1], h_map, w_map, c_map, s_map],
axis=1)
return geo_map, proposals
def detect_contours(self, tr_pred, tcl_pred, sin_pred, cos_pred, radii_pred):
# thresholding
tr_pred_mask = tr_pred > self.tr_thresh
tcl_pred_mask = tcl_pred > self.tcl_thresh
# multiply TR and TCL
tcl_mask = tcl_pred_mask * tr_pred_mask
# regularize
sin_pred, cos_pred = regularize_sin_cos(sin_pred, cos_pred)
# find disjoint regions
tcl_mask = fill_hole(tcl_mask)
tcl_contours, _ = cv2.findContours(tcl_mask.astype(np.uint8), cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
mask = np.zeros_like(tcl_mask)
bbox_contours = list()
for cont in tcl_contours:
deal_map = mask.copy()
cv2.drawContours(deal_map, [cont], -1, 1, -1)
if deal_map.sum() <= 100:
continue
text_map = tr_pred * deal_map
bboxs = self.bbox_transfor_inv(radii_pred, sin_pred, cos_pred, text_map, wclip=(4, 12))
# nms
boxes = lanms.merge_quadrangle_n9(bboxs.astype('float32'), 0.25)
boxes = boxes[:, :8].reshape((-1, 4, 2)).astype(np.int32)
boundary_point = None
if boxes.shape[0] > 1:
center = np.mean(boxes, axis=1).astype(np.int32).tolist()
paths, routes_path = minConnectPath(center)
boxes = boxes[routes_path]
top = np.mean(boxes[:, 0:2, :], axis=1).astype(np.int32).tolist()
bot = np.mean(boxes[:, 2:4, :], axis=1).astype(np.int32).tolist()
edge0 = self.select_edge(top + bot[::-1], boxes[0])
edge1 = self.select_edge(top + bot[::-1], boxes[-1])
if edge0 is not None:
top.insert(0, edge0[0])
bot.insert(0, edge0[1])
if edge1 is not None:
top.append(edge1[0])
bot.append(edge1[1])
boundary_point = np.array(top + bot[::-1])
elif boxes.shape[0] == 1:
top = boxes[0, 0:2, :].astype(np.int32).tolist()
bot = boxes[0, 2:4:-1, :].astype(np.int32).tolist()
boundary_point = np.array(top + bot)
if boundary_point is None:
continue
reconstruct_mask = mask.copy()
cv2.drawContours(reconstruct_mask, [boundary_point], -1, 1, -1)
if (reconstruct_mask * tr_pred_mask).sum() < reconstruct_mask.sum() * 0.5:
continue
# if reconstruct_mask.sum() < 200:
# continue
rect = cv2.minAreaRect(boundary_point)
if min(rect[1][0], rect[1][1]) < 10 or rect[1][0] * rect[1][1] < 300:
continue
bbox_contours.append([boundary_point, np.array(np.stack([top, bot], axis=1))])
return bbox_contours