update table recognition pipeline (#2047)

paddlex/inference/pipelines/table_recognition/table_recognition.py

@@ -17,7 +17,7 @@ from ..base import BasePipeline
 from ...predictors import create_predictor
 from ..ocr import OCRPipeline
 from ...components import CropByBoxes
-from ...results import TableResult, StructureTableResult
+from ...results import OCRResult, TableResult, StructureTableResult
 from .utils import *
 
 
@@ -66,17 +66,26 @@ class TableRecPipeline(BasePipeline):
                 single_img_res["img_path"] = layout_res["img_path"]
                 single_img_res["layout_result"] = layout_res
                 ocr_res = ocr_pred["result"]
-                single_img_res["ocr_result"] = ocr_res
                 all_subs_of_img = list(self._crop_by_boxes(layout_res))
                 # get cropped images with label 'table'
                 table_subs = []
                 for batch_subs in all_subs_of_img:
                     table_sub_list = []
                     for sub in batch_subs:
+                        box = sub["box"]
                         if sub["label"].lower() == "table":
                             table_sub_list.append(sub)
+                            _, ocr_res = self.get_ocr_result_by_bbox(box, ocr_res)
                     table_subs.append(table_sub_list)
-                single_img_res["table_result"] = self.get_table_result(table_subs)
+                table_res, all_table_ocr_res = self.get_table_result(table_subs)
+                for batch_table_ocr_res in all_table_ocr_res:
+                    for table_ocr_res in batch_table_ocr_res:
+                        ocr_res["dt_polys"].extend(table_ocr_res["dt_polys"])
+                        ocr_res["rec_text"].extend(table_ocr_res["rec_text"])
+                        ocr_res["rec_score"].extend(table_ocr_res["rec_score"])
+
+                single_img_res["table_result"] = table_res
+                single_img_res["ocr_result"] = OCRResult(ocr_res)
 
                 batch_structure_res.append({"result": TableResult(single_img_res)})
         yield batch_structure_res
@@ -84,24 +93,30 @@ class TableRecPipeline(BasePipeline):
     def get_ocr_result_by_bbox(self, box, ocr_res):
         dt_polys_list = []
         rec_text_list = []
-        unmatched_ocr_res = {"dt_polys": [], "rec_text": []}
-        for text_box, text_res in zip(ocr_res["dt_polys"], ocr_res["rec_text"]):
+        score_list = []
+        unmatched_ocr_res = {"dt_polys": [], "rec_text": [], "rec_score": []}
+        unmatched_ocr_res["img_path"] = ocr_res["img_path"]
+        for i, text_box in enumerate(ocr_res["dt_polys"]):
             text_box_area = convert_4point2rect(text_box)
-            if is_inside(box, text_box_area):
+            if is_inside(text_box_area, box):
                 dt_polys_list.append(text_box)
-                rec_text_list.append(text_res)
+                rec_text_list.append(ocr_res["rec_text"][i])
+                score_list.append(ocr_res["rec_score"][i])
             else:
                 unmatched_ocr_res["dt_polys"].append(text_box)
-                unmatched_ocr_res["rec_text"].append(text_res)
-        return (dt_polys_list, rec_text_list), unmatched_ocr_res
+                unmatched_ocr_res["rec_text"].append(ocr_res["rec_text"][i])
+                unmatched_ocr_res["rec_score"].append(ocr_res["rec_score"][i])
+        return (dt_polys_list, rec_text_list, score_list), unmatched_ocr_res
 
     def get_table_result(self, input_img):
         table_res_list = []
+        ocr_res_list = []
         table_index = 0
         for batch_input, batch_table_pred, batch_ocr_pred in zip(
             input_img, self.table_predictor(input_img), self.ocr_pipeline(input_img)
         ):
-            batch_res_list = []
+            batch_table_res = []
+            batch_ocr_res = []
             for input, table_pred, ocr_pred in zip(
                 batch_input, batch_table_pred, batch_ocr_pred
             ):
@@ -113,18 +128,24 @@ class TableRecPipeline(BasePipeline):
                     get_ori_coordinate_for_table(ori_x, ori_y, single_table_box),
                     dtype=np.float32,
                 )
+                ori_ocr_bbox_list = np.array(
+                    get_ori_coordinate_for_table(ori_x, ori_y, ocr_res["dt_polys"]),
+                    dtype=np.float32,
+                )
+                ocr_res["dt_polys"] = ori_ocr_bbox_list
                 html_res = self._match(single_table_res, ocr_res)
-                batch_res_list.append(
+                batch_table_res.append(
                     StructureTableResult(
                         {
                             "img_path": input["img_path"],
                             "bbox": ori_bbox_list,
                             "img_idx": table_index,
-                            "ocr_res": ocr_res,
                             "html": html_res,
                         }
                     )
                 )
+                batch_ocr_res.append(ocr_res)
                 table_index += 1
-            table_res_list.append(batch_res_list)
-        return table_res_list
+            table_res_list.append(batch_table_res)
+            ocr_res_list.append(batch_ocr_res)
+        return table_res_list, ocr_res_list

paddlex/inference/pipelines/table_recognition/utils.py

@@ -297,18 +297,11 @@ def get_ori_coordinate_for_table(x, y, table_bbox):
     Returns:
         list: list of original coordinates, eg. [[x1, y1, x2, y2, x3, y3, x4, y4]]
     """
-    bbox_list = []
-    for x1, y1, x2, y2, x3, y3, x4, y4 in table_bbox:
-        x1 = x + x1
-        y1 = y + y1
-        x2 = x + x2
-        y2 = y + y2
-        x3 = x + x3
-        y3 = y + y3
-        x4 = x + x4
-        y4 = y + y4
-        bbox_list.append([x1, y1, x2, y2, x3, y3, x4, y4])
-    return bbox_list
+    offset = np.array([x, y] * 4)
+    table_bbox = np.array(table_bbox)
+    if table_bbox.shape[-1] == 2:
+        offset = offset.reshape(4, 2)
+    return offset + table_bbox
 
 
 def is_inside(target_box, text_box):