diff --git a/ppocr/postprocess/db_postprocess.py b/ppocr/postprocess/db_postprocess.py index 27b428ef2e73c9abf81d3881b23979343c8595b2..1c42cd55cd8f85dff3df90e2f5365ccde8a725f3 100755 --- a/ppocr/postprocess/db_postprocess.py +++ b/ppocr/postprocess/db_postprocess.py @@ -38,6 +38,7 @@ class DBPostProcess(object): unclip_ratio=2.0, use_dilation=False, score_mode="fast", + visual_output=False, **kwargs): self.thresh = thresh self.box_thresh = box_thresh @@ -51,6 +52,7 @@ class DBPostProcess(object): self.dilation_kernel = None if not use_dilation else np.array( [[1, 1], [1, 1]]) + self.visual = visual_output def boxes_from_bitmap(self, pred, _bitmap, dest_width, dest_height): ''' @@ -169,12 +171,19 @@ class DBPostProcess(object): cv2.fillPoly(mask, contour.reshape(1, -1, 2).astype(np.int32), 1) return cv2.mean(bitmap[ymin:ymax + 1, xmin:xmax + 1], mask)[0] + def visual_output(self, pred): + im = np.array(pred[0] * 255).astype(np.uint8) + cv2.imwrite("db_probability_map.png", im) + print("The probalibity map is visualized in db_probability_map.png") + def __call__(self, outs_dict, shape_list): pred = outs_dict['maps'] if isinstance(pred, paddle.Tensor): pred = pred.numpy() pred = pred[:, 0, :, :] segmentation = pred > self.thresh + if self.visual: + self.visual_output(pred) boxes_batch = [] for batch_index in range(pred.shape[0]): diff --git a/tools/infer/utility.py b/tools/infer/utility.py index 4f4caabaf89a87a1c5d17b2f48e7dd62d9f88ebf..75a8b2345d50f09d6548ca6f1dac32105a707ee5 100644 --- a/tools/infer/utility.py +++ b/tools/infer/utility.py @@ -55,6 +55,7 @@ def init_args(): parser.add_argument("--max_batch_size", type=int, default=10) parser.add_argument("--use_dilation", type=str2bool, default=False) parser.add_argument("--det_db_score_mode", type=str, default="fast") + parser.add_argument("--visual_output", type=str2bool, default=False) # EAST parmas parser.add_argument("--det_east_score_thresh", type=float, default=0.8) parser.add_argument("--det_east_cover_thresh", type=float, default=0.1)