浏览代码

Add new table related methods (#3296)

* fix bugs

* refine codes

* fix bugs

* refine code

* add new algorithm

* refine codes

* refine codes

* refine table method
Liu Jiaxuan 9 月之前
父节点
当前提交
d35b1f1c28
共有 1 个文件被更改,包括 45 次插入41 次删除
  1. 45 41
      paddlex/inference/pipelines/table_recognition/pipeline_v2.py

+ 45 - 41
paddlex/inference/pipelines/table_recognition/pipeline_v2.py

@@ -338,17 +338,15 @@ class TableRecognitionPipelineV2(BasePipeline):
         Returns:
         Returns:
             list of list: List of adjusted bounding boxes relative to table_box, for boxes fully inside table_box.
             list of list: List of adjusted bounding boxes relative to table_box, for boxes fully inside table_box.
         """
         """
-        tol = 5
+        tol=0
         # Extract coordinates from table_box
         # Extract coordinates from table_box
         x_min_t, y_min_t, x_max_t, y_max_t = table_box
         x_min_t, y_min_t, x_max_t, y_max_t = table_box
-
         adjusted_boxes = []
         adjusted_boxes = []
         for box in ocr_det_boxes:
         for box in ocr_det_boxes:
             x_min_b, y_min_b, x_max_b, y_max_b = box
             x_min_b, y_min_b, x_max_b, y_max_b = box
-
             # Check if the box is fully inside table_box
             # Check if the box is fully inside table_box
             if (x_min_b+tol >= x_min_t and y_min_b+tol >= y_min_t and
             if (x_min_b+tol >= x_min_t and y_min_b+tol >= y_min_t and
-                x_max_b+tol <= x_max_t and y_max_b+tol <= y_max_t):
+                x_max_b-tol <= x_max_t and y_max_b-tol <= y_max_t):
                 # Adjust the coordinates to be relative to table_box
                 # Adjust the coordinates to be relative to table_box
                 adjusted_box = [
                 adjusted_box = [
                     x_min_b - x_min_t,  # Adjust x1
                     x_min_b - x_min_t,  # Adjust x1
@@ -398,7 +396,7 @@ class TableRecognitionPipelineV2(BasePipeline):
             box1_area = (box1[2] - box1[0]) * (box1[3] - box1[1])
             box1_area = (box1[2] - box1[0]) * (box1[3] - box1[1])
             box2_area = (box2[2] - box2[0]) * (box2[3] - box2[1])
             box2_area = (box2[2] - box2[0]) * (box2[3] - box2[1])
             # Calculate the IoU
             # Calculate the IoU
-            iou = intersection_area / float(box1_area + box2_area - intersection_area)
+            iou = intersection_area / float(box1_area)
             return iou
             return iou
 
 
         # Function to combine rectangles into N rectangles
         # Function to combine rectangles into N rectangles
@@ -407,11 +405,11 @@ class TableRecognitionPipelineV2(BasePipeline):
             Combine rectangles into N rectangles based on geometric proximity.
             Combine rectangles into N rectangles based on geometric proximity.
 
 
             Args:
             Args:
-            rectangles (list of list of int): A list of rectangles, each represented by [x1, y1, x2, y2].
-            N (int): The desired number of combined rectangles.
+                rectangles (list of list of int): A list of rectangles, each represented by [x1, y1, x2, y2].
+                N (int): The desired number of combined rectangles.
 
 
             Returns:
             Returns:
-            list of list of int: A list of N combined rectangles.
+                list of list of int: A list of N combined rectangles.
             """
             """
             # Number of input rectangles
             # Number of input rectangles
             num_rects = len(rectangles)
             num_rects = len(rectangles)
@@ -453,44 +451,52 @@ class TableRecognitionPipelineV2(BasePipeline):
         cells_det_results = np.array(cells_det_results)
         cells_det_results = np.array(cells_det_results)
         cells_det_scores = np.array(cells_det_scores)
         cells_det_scores = np.array(cells_det_scores)
         ocr_det_results = np.array(ocr_det_results)
         ocr_det_results = np.array(ocr_det_results)
+        more_cells_flag = False
         if len(cells_det_results) == html_pred_boxes_nums:
         if len(cells_det_results) == html_pred_boxes_nums:
             return cells_det_results
             return cells_det_results
         # Step 1: If cells_det_results has more rectangles than html_pred_boxes_nums
         # Step 1: If cells_det_results has more rectangles than html_pred_boxes_nums
         elif len(cells_det_results) > html_pred_boxes_nums:
         elif len(cells_det_results) > html_pred_boxes_nums:
-                return combine_rectangles(cells_det_results, html_pred_boxes_nums)
+            more_cells_flag = True
+            # Select the indices of the top html_pred_boxes_nums scores
+            top_indices = np.argsort(-cells_det_scores)[:html_pred_boxes_nums]
+            # Adjust the corresponding rectangles
+            cells_det_results = cells_det_results[top_indices].tolist()
+        # Threshold for IoU
+        iou_threshold = 0.6
+        # List to store ocr_miss_boxes
+        ocr_miss_boxes = []
+        # For each rectangle in ocr_det_results
+        for ocr_rect in ocr_det_results:
+            merge_ocr_box_iou = []
+            # Flag to indicate if ocr_rect has IoU >= threshold with any cell_rect
+            has_large_iou = False
+            # For each rectangle in cells_det_results
+            for cell_rect in cells_det_results:
+                # Compute IoU
+                iou = compute_iou(ocr_rect, cell_rect)
+                if iou > 0:
+                    merge_ocr_box_iou.append(iou)
+                if (iou>=iou_threshold) or (sum(merge_ocr_box_iou)>=iou_threshold):
+                    has_large_iou = True
+                    break
+            if not has_large_iou:
+                ocr_miss_boxes.append(ocr_rect)
+        # If no ocr_miss_boxes, return cells_det_results
+        if len(ocr_miss_boxes) == 0:
+            final_results = cells_det_results if more_cells_flag==True else cells_det_results.tolist()
         else:
         else:
-            # return cells_det_results
-            # Threshold for IoU
-            iou_threshold = 0.1
-            # List to store ocr_miss_boxes
-            ocr_miss_boxes = []
-            # For each rectangle in ocr_det_results
-            for ocr_rect in ocr_det_results:
-                # Flag to indicate if ocr_rect has IoU >= threshold with any cell_rect
-                has_large_iou = False
-                # For each rectangle in cells_det_results
-                for cell_rect in cells_det_results:
-                    # Compute IoU
-                    iou = compute_iou(ocr_rect, cell_rect)
-                    if iou >= iou_threshold:
-                        has_large_iou = True
-                        break
-                if not has_large_iou:
-                    ocr_miss_boxes.append(ocr_rect)
-            # If no ocr_miss_boxes, return cells_det_results
-            if len(ocr_miss_boxes) == 0:
-                return cells_det_results.tolist()
+            if more_cells_flag == True:
+                final_results = combine_rectangles(cells_det_results+ocr_miss_boxes, html_pred_boxes_nums)
             else:
             else:
                 # Need to combine ocr_miss_boxes into N rectangles
                 # Need to combine ocr_miss_boxes into N rectangles
                 N = html_pred_boxes_nums - len(cells_det_results)
                 N = html_pred_boxes_nums - len(cells_det_results)
-                if len(ocr_miss_boxes) == N:
-                    return cells_det_results.tolist() + ocr_miss_boxes
-                else:
-                    # Combine ocr_miss_boxes into N rectangles
-                    ocr_supp_boxes = combine_rectangles(ocr_miss_boxes, N)
-                    # Combine cells_det_results and ocr_supp_boxes
-                    final_results = np.concatenate((cells_det_results, ocr_supp_boxes), axis=0)
-                    return final_results.tolist()
+                # Combine ocr_miss_boxes into N rectangles
+                ocr_supp_boxes = combine_rectangles(ocr_miss_boxes, N)
+                # Combine cells_det_results and ocr_supp_boxes
+                final_results = np.concatenate((cells_det_results, ocr_supp_boxes), axis=0).tolist()
+        if len(final_results) <= 0.6*html_pred_boxes_nums:
+            final_results = combine_rectangles(ocr_det_results, html_pred_boxes_nums)
+        return final_results
 
 
     def predict_single_table_recognition_res(
     def predict_single_table_recognition_res(
         self,
         self,
@@ -525,9 +531,7 @@ class TableRecognitionPipelineV2(BasePipeline):
                 self.wireless_table_cells_detection_model(image_array, threshold=0.3)
                 self.wireless_table_cells_detection_model(image_array, threshold=0.3)
             ) # Setting the threshold to 0.3 can improve the accuracy of table cells detection. 
             ) # Setting the threshold to 0.3 can improve the accuracy of table cells detection. 
               # If you really want more or fewer table cells detection boxes, the threshold can be adjusted.
               # If you really want more or fewer table cells detection boxes, the threshold can be adjusted.
-        table_structure_result = self.extract_results(
-            table_structure_pred, "table_stru"
-        )
+        table_structure_result = self.extract_results(table_structure_pred, "table_stru")
         table_cells_result, table_cells_score = self.extract_results(table_cells_pred, "det")
         table_cells_result, table_cells_score = self.extract_results(table_cells_pred, "det")
         table_cells_result, table_cells_score = self.cells_det_results_nms(table_cells_result, table_cells_score)
         table_cells_result, table_cells_score = self.cells_det_results_nms(table_cells_result, table_cells_score)
         ocr_det_boxes = self.get_region_ocr_det_boxes(overall_ocr_res["rec_boxes"].tolist(), table_box)
         ocr_det_boxes = self.get_region_ocr_det_boxes(overall_ocr_res["rec_boxes"].tolist(), table_box)