forgetest5
/
PaddleOCR
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

change DBHead output to dict

This commit is contained in:
WenmuZhou 2020-12-15 23:49:50 +08:00
parent 57e6edd97c
commit d3b609ee09
4 changed files with 10 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
ppocr/losses/det_db_loss.py
View File

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

4
ppocr/modeling/heads/det_db_head.py
View File

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

6
ppocr/postprocess/db_postprocess.py
View File

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

2
tools/infer/predict_det.py
View File

@ -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']