Merge pull request #6670 from LDOUBLEV/2.5

add visual output for db postprocess

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

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]]}......
+```
+
+**注：如果您希望在自己的数据集上训练，建议按照上述数据个数准备数据集。**
+
 执行预测：
 
 ```

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

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)