modify ppocr2example to fix bugs caused by empty ocr input

PaddlePaddle · May 19, 2023 · 722e1d9 · 722e1d9
1 parent ea72e05
commit 722e1d9
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 27 deletions.
diff --git a/model_zoo/ernie-layout/deploy/python/predictor.py b/model_zoo/ernie-layout/deploy/python/predictor.py
@@ -754,21 +754,12 @@ def infer(self, data):
 
     def predict(self, docs):
         input_data = []
-        invalid_idx, input_idx, idx_map = set(), [], {}
-        for idx, doc in enumerate(docs):
+        for doc in docs:
             ocr_result = self.ocr.ocr(doc, cls=True)
             # Compatible with paddleocr>=2.6.0.2
             ocr_result = ocr_result[0] if len(ocr_result) == 1 else ocr_result
-            if ocr_result:
-                # Only process images with ocr results
-                example = ppocr2example(ocr_result, doc)
-                input_data.append(example)
-                input_idx.append(idx)
-            else:
-                invalid_idx.add(idx)
-
-        # Save maping between original index in docs and new index(index in valid inputs)
-        idx_map = {j: i for i, j in enumerate(input_idx)}
+            example = ppocr2example(ocr_result, doc)
+            input_data.append(example)
 
         inputs = collections.defaultdict(list)
         for data in input_data:
@@ -789,20 +780,8 @@ def predict(self, docs):
                 preds[0].extend(output[0].tolist())
                 preds[1].extend(output[1].tolist())
         results = self.postprocess(preds)
-
-        # Merge results of invalid docs & valid docs input
-        merged_results = []
-        for i in range(len(docs)):
-            if i in invalid_idx:
-                # Set empty result for invalid doc
-                result = ""
-            else:
-                result = results[idx_map[i]]
-            merged_results.append(result)
-
         formatted_results = []
-        # for doc, res in zip(docs, results):
-        for doc, res in zip(docs, merged_results):
+        for doc, res in zip(docs, results):
             formatted_result = {"doc": doc, "result": res}
             formatted_results.append(formatted_result)
         return formatted_results

diff --git a/paddlenlp/utils/image_utils.py b/paddlenlp/utils/image_utils.py
@@ -695,8 +695,8 @@ def ppocr2example(ocr_res, img_path):
     doc_tokens = []
     doc_boxes = []
 
-    im_w_box = max([seg["bbox"].left + seg["bbox"].width for seg in segments]) + 20
-    im_h_box = max([seg["bbox"].top + seg["bbox"].height for seg in segments]) + 20
+    im_w_box = max([seg["bbox"].left + seg["bbox"].width for seg in segments]) + 20 if segments else 0
+    im_h_box = max([seg["bbox"].top + seg["bbox"].height for seg in segments]) + 20 if segments else 0
     img = Image.open(img_path)
     im_w, im_h = img.size
     im_w, im_h = max(im_w, im_w_box), max(im_h, im_h_box)