提交 d3b609ee 编写于 作者: W WenmuZhou

change DBHead output to dict

上级 57e6edd9
...@@ -47,6 +47,7 @@ class DBLoss(nn.Layer): ...@@ -47,6 +47,7 @@ class DBLoss(nn.Layer):
negative_ratio=ohem_ratio) negative_ratio=ohem_ratio)
def forward(self, predicts, labels): def forward(self, predicts, labels):
predicts = predicts['maps']
label_threshold_map, label_threshold_mask, label_shrink_map, label_shrink_mask = labels[ label_threshold_map, label_threshold_mask, label_shrink_map, label_shrink_mask = labels[
1:] 1:]
shrink_maps = predicts[:, 0, :, :] shrink_maps = predicts[:, 0, :, :]
......
...@@ -120,9 +120,9 @@ class DBHead(nn.Layer): ...@@ -120,9 +120,9 @@ class DBHead(nn.Layer):
def forward(self, x): def forward(self, x):
shrink_maps = self.binarize(x) shrink_maps = self.binarize(x)
if not self.training: if not self.training:
return shrink_maps return {'maps': shrink_maps}
threshold_maps = self.thresh(x) threshold_maps = self.thresh(x)
binary_maps = self.step_function(shrink_maps, threshold_maps) binary_maps = self.step_function(shrink_maps, threshold_maps)
y = paddle.concat([shrink_maps, threshold_maps, binary_maps], axis=1) y = paddle.concat([shrink_maps, threshold_maps, binary_maps], axis=1)
return y return {'maps': y}
...@@ -40,7 +40,8 @@ class DBPostProcess(object): ...@@ -40,7 +40,8 @@ class DBPostProcess(object):
self.max_candidates = max_candidates self.max_candidates = max_candidates
self.unclip_ratio = unclip_ratio self.unclip_ratio = unclip_ratio
self.min_size = 3 self.min_size = 3
self.dilation_kernel = None if not use_dilation else np.array([[1, 1], [1, 1]]) self.dilation_kernel = None if not use_dilation else np.array(
[[1, 1], [1, 1]])
def boxes_from_bitmap(self, pred, _bitmap, dest_width, dest_height): def boxes_from_bitmap(self, pred, _bitmap, dest_width, dest_height):
''' '''
...@@ -132,7 +133,8 @@ class DBPostProcess(object): ...@@ -132,7 +133,8 @@ class DBPostProcess(object):
cv2.fillPoly(mask, box.reshape(1, -1, 2).astype(np.int32), 1) cv2.fillPoly(mask, box.reshape(1, -1, 2).astype(np.int32), 1)
return cv2.mean(bitmap[ymin:ymax + 1, xmin:xmax + 1], mask)[0] return cv2.mean(bitmap[ymin:ymax + 1, xmin:xmax + 1], mask)[0]
def __call__(self, pred, shape_list): def __call__(self, outs_dict, shape_list):
pred = outs_dict['maps']
if isinstance(pred, paddle.Tensor): if isinstance(pred, paddle.Tensor):
pred = pred.numpy() pred = pred.numpy()
pred = pred[:, 0, :, :] pred = pred[:, 0, :, :]
......
...@@ -178,7 +178,7 @@ class TextDetector(object): ...@@ -178,7 +178,7 @@ class TextDetector(object):
preds['f_tco'] = outputs[2] preds['f_tco'] = outputs[2]
preds['f_tvo'] = outputs[3] preds['f_tvo'] = outputs[3]
else: else:
preds = outputs[0] preds['maps'] = outputs[0]
post_result = self.postprocess_op(preds, shape_list) post_result = self.postprocess_op(preds, shape_list)
dt_boxes = post_result[0]['points'] dt_boxes = post_result[0]['points']
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册