def detect_contours(self, image, tr_pred, tcl_pred, sin_pred, cos_pred,
radii_pred):
"""
Input: FCN output, Output: text detection after post-processing
:param image: (np.array) input image (3, H, W)
:param tr_pred: (np.array), text region prediction, (2, H, W)
:param tcl_pred: (np.array), text center line prediction, (2, H, W)
:param sin_pred: (np.array), sin prediction, (H, W)
:param cos_pred: (np.array), cos line prediction, (H, W)
:param radii_pred: (np.array), radii prediction, (H, W)
:return:
(list), tcl array: (n, 3), 3 denotes (x, y, radii)
"""
# thresholding
tr_pred_mask = tr_pred[1] > self.tr_thresh
tcl_pred_mask = tcl_pred[1] > 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 tcl in each predicted mask
detect_result = self.build_tcl(tcl_mask, sin_pred, cos_pred,
radii_pred)
return self.postprocessing(image, detect_result, tr_pred_mask)
def detect(self, tr_pred, tcl_pred, sin_pred, cos_pred, radii_pred):
# thresholding
tr_pred_mask = tr_pred[1] > self.tr_thresh
tcl_pred_mask = tcl_pred[1] > self.tcl_thresh
# multiply TR and TCL
tcl = tcl_pred_mask * tr_pred_mask
# regularize
sin_pred, cos_pred = regularize_sin_cos(sin_pred, cos_pred)
# find tcl in each predicted mask
detect_result = self.build_tcl(tcl, sin_pred, cos_pred, radii_pred)
return detect_result
for idx in range(0, len(trainset)):
t0 = time.time()
img, train_mask, tr_mask, tcl_mask, radius_map, sin_map, cos_map, gt_roi = trainset[
idx]
img, train_mask, tr_mask, tcl_mask, radius_map, sin_map, cos_map, gt_roi \
= map(lambda x: x.cpu().numpy(), (img, train_mask, tr_mask, tcl_mask, radius_map, sin_map, cos_map, gt_roi))
img = img.transpose(1, 2, 0)
img = ((img * stds + means) * 255).astype(np.uint8)
print(idx, img.shape)
top_map = radius_map[:, :, 0]
bot_map = radius_map[:, :, 1]
print(radius_map.shape)
sin_map, cos_map = regularize_sin_cos(sin_map, cos_map)
ret, labels = cv2.connectedComponents(tcl_mask[:, :,
0].astype(np.uint8),
connectivity=8)
cv2.imshow(
"labels0",
cav.heatmap(np.array(labels * 255 / np.max(labels),
dtype=np.uint8)))
print(np.sum(tcl_mask[:, :, 1]))
t0 = time.time()
for bbox_idx in range(1, ret):
bbox_mask = labels == bbox_idx
text_map = tcl_mask[:, :, 0] * bbox_mask
boxes = bbox_transfor_inv(radius_map,
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
