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/ppstructure/docs/kie.md b/ppstructure/docs/kie.md index 35498b33478d1010fd2548dfcb8586b4710723a1..8fd5a7921e67922b69c9da1f72f7bb514c95323a 100644 --- a/ppstructure/docs/kie.md +++ b/ppstructure/docs/kie.md @@ -19,6 +19,24 @@ SDMGR是一个关键信息提取算法,将每个检测到的文本区域分类 wget https://paddleocr.bj.bcebos.com/dygraph_v2.1/kie/wildreceipt.tar && tar xf wildreceipt.tar ``` +数据集格式: +``` +./wildreceipt +├── class_list.txt # box内的文本类别,比如金额、时间、日期等。 +├── dict.txt # 识别的字典文件,数据集中包含的字符列表 +├── wildreceipt_train.txt # 训练数据标签文件 +└── wildreceipt_test.txt # 评估数据标签文件 +└── image_files/ # 图像数据文件夹 +``` + +其中标签文件里的格式为: +``` +" 图像文件名 json.dumps编码的图像标注信息" +image_files/Image_16/11/d5de7f2a20751e50b84c747c17a24cd98bed3554.jpeg [{"label": 1, "transcription": "SAFEWAY", "points": [[550.0, 190.0], [937.0, 190.0], [937.0, 104.0], [550.0, 104.0]]}, {"label": 25, "transcription": "TM", "points": [[1048.0, 211.0], [1074.0, 211.0], [1074.0, 196.0], [1048.0, 196.0]]}, {"label": 25, "transcription": "ATOREMGRTOMMILAZZO", "points": [[535.0, 239.0], [833.0, 239.0], [833.0, 200.0], [535.0, 200.0]]}, {"label": 5, "transcription": "703-777-5833", "points": [[907.0, 256.0], [1081.0, 256.0], [1081.0, 223.0], [907.0, 223.0]]}...... +``` + +**注:如果您希望在自己的数据集上训练,建议按照上述数据个数准备数据集。** + 执行预测: ``` diff --git a/ppstructure/docs/kie_en.md b/ppstructure/docs/kie_en.md index 1fe38b0b399e9290526dafa5409673dc87026db7..e895ee88d65911f4151096f56c17c9c13af3277c 100644 --- a/ppstructure/docs/kie_en.md +++ b/ppstructure/docs/kie_en.md @@ -18,6 +18,22 @@ This section provides a tutorial example on how to quickly use, train, and evalu wget https://paddleocr.bj.bcebos.com/dygraph_v2.1/kie/wildreceipt.tar && tar xf wildreceipt.tar ``` +The dataset format are as follows: +``` +./wildreceipt +├── class_list.txt # The text category inside the box, such as amount, time, date, etc. +├── dict.txt # A recognized dictionary file, a list of characters contained in the dataset +├── wildreceipt_train.txt # training data label file +└── wildreceipt_test.txt # testing data label file +└── image_files/ # image dataset file +``` + +The format in the label file is: +``` +" The image file path Image annotation information encoded by json.dumps" +image_files/Image_16/11/d5de7f2a20751e50b84c747c17a24cd98bed3554.jpeg [{"label": 1, "transcription": "SAFEWAY", "points": [[550.0, 190.0], [937.0, 190.0], [937.0, 104.0], [550.0, 104.0]]}, {"label": 25, "transcription": "TM", "points": [[1048.0, 211.0], [1074.0, 211.0], [1074.0, 196.0], [1048.0, 196.0]]}, {"label": 25, "transcription": "ATOREMGRTOMMILAZZO", "points": [[535.0, 239.0], [833.0, 239.0], [833.0, 200.0], [535.0, 200.0]]}, {"label": 5, "transcription": "703-777-5833", "points": [[907.0, 256.0], [1081.0, 256.0], [1081.0, 223.0], [907.0, 223.0]]}...... +``` + Download the pretrained model and predict the result: ```shell diff --git a/tools/infer/utility.py b/tools/infer/utility.py index 4f4caabaf89a87a1c5d17b2f48e7dd62d9f88ebf..6d9935a70e79bb20c5f6380783911ef141b0be17 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("--vis_seg_map", 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)