From 026f96b6578dde670df21ec0209039529834d616 Mon Sep 17 00:00:00 2001 From: WenmuZhou Date: Thu, 29 Jul 2021 19:33:47 +0800 Subject: [PATCH] add export model of pse --- tools/infer/predict_det.py | 4 +--- tools/infer/utility.py | 2 +- 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/tools/infer/predict_det.py b/tools/infer/predict_det.py index f727b4f5..728c3667 100755 --- a/tools/infer/predict_det.py +++ b/tools/infer/predict_det.py @@ -223,9 +223,7 @@ class TextDetector(object): #self.predictor.try_shrink_memory() post_result = self.postprocess_op(preds, shape_list) dt_boxes = post_result[0]['points'] - if self.det_algorithm == "SAST" and self.det_sast_polygon: - dt_boxes = self.filter_tag_det_res_only_clip(dt_boxes, ori_im.shape) - elif self.det_algorithm == "PSE" and self.det_pse_box_type=='poly': + if (self.det_algorithm == "SAST" and self.det_sast_polygon) or (self.det_algorithm == "PSE" and self.det_pse_box_type=='poly'): dt_boxes = self.filter_tag_det_res_only_clip(dt_boxes, ori_im.shape) else: dt_boxes = self.filter_tag_det_res(dt_boxes, ori_im.shape) diff --git a/tools/infer/utility.py b/tools/infer/utility.py index 9bb6136e..4ac48cf7 100755 --- a/tools/infer/utility.py +++ b/tools/infer/utility.py @@ -66,7 +66,7 @@ def init_args(): parser.add_argument("--det_pse_thresh", type=float, default=0) parser.add_argument("--det_pse_box_thresh", type=float, default=0.85) parser.add_argument("--det_pse_min_area", type=float, default=16) - parser.add_argument("--det_pse_box_type", type=str, default='poly') + parser.add_argument("--det_pse_box_type", type=str, default='box') parser.add_argument("--det_pse_scale", type=int, default=1) # params for text recognizer -- GitLab