Merge pull request #119 from deeperrrr/batch_process

feat: 增加slanet-plus表格结构模型批量推理功能

Author: SWHL

batch_demo.py

@@ -0,0 +1,66 @@
+# -*- encoding: utf-8 -*-
+# @Author: deeperrrr
+# @Contact: 3545615231@qq.com
+import cv2
+import numpy as np
+from pathlib import Path
+from typing import List
+from tqdm import tqdm
+
+from rapidocr import EngineType, RapidOCR
+from rapid_table import ModelType, RapidTable, RapidTableInput
+
+ocr_engine = RapidOCR(
+    params={
+        "Det.engine_type": EngineType.TORCH,
+        "Cls.engine_type": EngineType.TORCH,
+        "Rec.engine_type": EngineType.TORCH,
+    }
+)
+img_dir_path = "/data/images"  # 图片文件夹
+ocr_results = []
+batch_size = 4
+
+# input_args = RapidTableInput(model_type=ModelType.UNITABLE)
+input_args = RapidTableInput(model_type=ModelType.SLANETPLUS)
+table_engine = RapidTable(input_args)
+
+
+def load_images_original_size(img_dir: str) -> List[np.ndarray]:
+    img_dir = Path(img_dir)
+    if not img_dir.exists():
+        raise FileNotFoundError(f"目录不存在: {img_dir}")
+
+    image_paths = []
+    for ext in ['*.jpg', '*.jpeg', '*.png']:
+        image_paths.extend(list(img_dir.glob(ext)))
+
+    images = []
+    for img_path in tqdm(image_paths, desc="加载图像"):
+        img = cv2.imread(str(img_path))
+        images.append(img)
+    return images
+
+
+def dynamic_batch_process(table_engine: RapidTable, images: List[np.ndarray], ocr_results: List[List],
+                          batch_size: int = 1):
+    all_results = []
+    for i in tqdm(range(0, len(images), batch_size), desc=f"表格批量推理, batch_size={batch_size}"):
+        batch_imgs = images[i:i + batch_size]
+        batch_ocrs = ocr_results[i:i + batch_size]
+        results = table_engine(batch_imgs, batch_ocrs, batch_size)
+        all_results.extend(results)
+    return all_results
+
+
+images = load_images_original_size(img_dir_path)
+for img in tqdm(images, desc="OCR处理"):
+    ori_ocr_res = ocr_engine(img)
+    ocr_result = [ori_ocr_res.boxes, ori_ocr_res.txts, ori_ocr_res.scores]
+    ocr_results.append(ocr_result)
+
+# 批量表格结构识别
+results = dynamic_batch_process(table_engine, images, ocr_results, batch_size)  # batch_size默认4
+
+for i, result in enumerate(results):
+    result.vis(save_dir="outputs", save_name=f"vis_{i}")

rapid_table/main.py

@@ -64,27 +64,71 @@ def _init_table_structer(self):
 
         return PPTableStructurer(asdict(self.cfg))
 
+    def _batch_process(
+            self,
+            img_contents: List[Union[str, np.ndarray, bytes, Path]],
+            ocr_results: Optional[List] = None,
+            batch_size: int = 4,
+    ) -> List[RapidTableOutput]:
+        """批量处理图像"""
+        s = time.perf_counter()
+
+        images = []
+        for img_content in img_contents:
+            img = self.load_img(img_content)
+            images.append(img)
+
+        batch_dt_boxes = []
+        batch_rec_res = []
+
+        for i, img in enumerate(images):
+            dt_boxes, rec_res = get_boxes_recs(ocr_results[i], img.shape[:2])
+            batch_dt_boxes.append(dt_boxes)
+            batch_rec_res.append(rec_res)
+
+        # 批量表格结构识别
+        batch_results = self.table_structure(images, batch_size)
+
+        output_results = []
+        for i, (img, (pred_structures, cell_bboxes, _)) in enumerate(zip(images, batch_results)):
+            logic_points = self.table_matcher.decode_logic_points(pred_structures)
+            pred_html = self.get_table_matcher(
+                pred_structures, cell_bboxes, batch_dt_boxes[i], batch_rec_res[i]
+            )
+            result = RapidTableOutput(img, pred_html, cell_bboxes, logic_points, 0)
+            output_results.append(result)
+
+        total_elapse = time.perf_counter() - s
+        for result in output_results:
+            result.elapse = total_elapse / len(output_results)
+
+        return output_results
+
     def __call__(
-        self,
-        img_content: Union[str, np.ndarray, bytes, Path],
-        ocr_results: Optional[Tuple[np.ndarray, Tuple[str], Tuple[float]]] = None,
+            self,
+            img_content: Union[str, np.ndarray, bytes, Path],
+            ocr_results: Optional[Tuple[np.ndarray, Tuple[str], Tuple[float]]] = None,
+            batch_size: int = 1,
     ) -> RapidTableOutput:
-        s = time.perf_counter()
+        if batch_size > 1:
+            return self._batch_process(img_content, ocr_results)
+        else:
+            s = time.perf_counter()
 
-        img = self.load_img(img_content)
+            img = self.load_img(img_content)
 
-        dt_boxes, rec_res = self.get_ocr_results(img, ocr_results)
-        pred_structures, cell_bboxes, logic_points = self.get_table_rec_results(img)
+            dt_boxes, rec_res = self.get_ocr_results(img, ocr_results)
+            pred_structures, cell_bboxes, logic_points = self.get_table_rec_results(img)
 
-        pred_html = self.get_table_matcher(
-            pred_structures, cell_bboxes, dt_boxes, rec_res
-        )
+            pred_html = self.get_table_matcher(
+                pred_structures, cell_bboxes, dt_boxes, rec_res
+            )
 
-        elapse = time.perf_counter() - s
-        return RapidTableOutput(img, pred_html, cell_bboxes, logic_points, elapse)
+            elapse = time.perf_counter() - s
+            return RapidTableOutput(img, pred_html, cell_bboxes, logic_points, elapse)
 
     def get_ocr_results(
-        self, img: np.ndarray, ocr_results: Tuple[np.ndarray, Tuple[str], Tuple[float]]
+            self, img: np.ndarray, ocr_results: Tuple[np.ndarray, Tuple[str], Tuple[float]]
     ) -> Tuple[Optional[np.ndarray], Optional[np.ndarray]]:
         if ocr_results is not None:
             return get_boxes_recs(ocr_results, img.shape[:2])

rapid_table/table_structure/pp_structure/main.py

@@ -21,7 +21,7 @@
 from ...inference_engine.base import get_engine
 from ..utils import get_struct_str
 from .post_process import TableLabelDecode
-from .pre_process import TablePreprocess
+from .pre_process import TablePreprocess, BatchTablePreprocess
 
 
 class PPTableStructurer:
@@ -32,11 +32,15 @@ def __init__(self, cfg: Dict[str, Any]):
         self.cfg = cfg
 
         self.preprocess_op = TablePreprocess()
+        self.batch_preprocess_op = BatchTablePreprocess()
 
         self.character = self.session.get_character_list()
         self.postprocess_op = TableLabelDecode(self.character)
 
-    def __call__(self, ori_img: np.ndarray) -> Tuple[List[str], np.ndarray, float]:
+    def __call__(self, ori_img: np.ndarray, batch_size: int = 4) -> Tuple[List[str], np.ndarray, float]:
+        if batch_size > 1:
+            return self.batch_process(ori_img)
+
         s = time.perf_counter()
 
         img, shape_list = self.preprocess_op(ori_img)
@@ -55,8 +59,52 @@ def __call__(self, ori_img: np.ndarray) -> Tuple[List[str], np.ndarray, float]:
         elapse = time.perf_counter() - s
         return table_struct_str, cell_bboxes, elapse
 
+    def batch_process(self, img_list: List[np.ndarray]) -> List[Tuple[List[str], np.ndarray, float]]:
+        """批量处理图像列表
+        Args:
+            img_list: 图像列表
+
+        Returns:
+            结果列表，每个元素包含 (table_struct_str, cell_bboxes, elapse)
+        """
+        starttime = time.perf_counter()
+
+        batch_data = self.batch_preprocess_op(img_list)
+
+        preprocessed_images = batch_data[0]
+        shape_lists = batch_data[1]
+        preprocessed_images = np.array(preprocessed_images)
+
+        bbox_preds, struct_probs = self.session(preprocessed_images)
+
+        batch_size = preprocessed_images.shape[0]
+        results = []
+
+        for i in range(batch_size):
+            single_bbox_preds = bbox_preds[i:i + 1]
+            single_struct_probs = struct_probs[i:i + 1]
+            single_shape_list = np.array([shape_lists[i]])
+
+            post_result = self.postprocess_op(single_bbox_preds, single_struct_probs, [single_shape_list])
+
+            table_struct_str = get_struct_str(post_result["structure_batch_list"][0][0])
+            cell_bboxes = post_result["bbox_batch_list"][0]
+
+            if self.cfg["model_type"] == ModelType.SLANETPLUS:
+                cell_bboxes = self.rescale_cell_bboxes(img_list[i], cell_bboxes)
+
+            cell_bboxes = self.filter_blank_bbox(cell_bboxes)
+
+            results.append((table_struct_str, cell_bboxes, 0))
+
+        total_elapse = time.perf_counter() - starttime
+        for i in range(len(results)):
+            results[i] = (results[i][0], results[i][1], total_elapse / batch_size)
+
+        return results
+
     def rescale_cell_bboxes(
-        self, img: np.ndarray, cell_bboxes: np.ndarray
+            self, img: np.ndarray, cell_bboxes: np.ndarray
     ) -> np.ndarray:
         h, w = img.shape[:2]
         resized = 488

rapid_table/table_structure/pp_structure/pre_process.py

@@ -37,7 +37,7 @@ def normalize(self, img: np.ndarray) -> np.ndarray:
         return (img.astype("float32") * self.scale - self.mean) / self.std
 
     def pad_img(
-        self, img: np.ndarray, shape: List[float]
+            self, img: np.ndarray, shape: List[float]
     ) -> Tuple[np.ndarray, List[float]]:
         padding_img = np.zeros((self.max_len, self.max_len, 3), dtype=np.float32)
         h, w = img.shape[:2]
@@ -47,3 +47,39 @@ def pad_img(
 
     def to_chw(self, img: np.ndarray) -> np.ndarray:
         return img.transpose((2, 0, 1))
+
+
+class BatchTablePreprocess:
+    """批量表格预处理类"""
+
+    def __init__(self):
+        self.preprocess = TablePreprocess()
+
+    def __call__(self, img_list: List[np.ndarray]) -> Tuple[List[np.ndarray], List[List[float]]]:
+        """批量处理图像
+
+        Args:
+            img_list: 图像列表
+
+        Returns:
+            预处理后的图像列表和形状信息列表
+        """
+        if not img_list:
+            return None, None
+
+        processed_imgs = []
+        shape_lists = []
+
+        for img in img_list:
+            if img is None:
+                continue
+
+            img_processed, shape_list = self.preprocess.resize_image(img)
+            img_processed = self.preprocess.normalize(img_processed)
+            img_processed, shape_list = self.preprocess.pad_img(img_processed, shape_list)
+            img_processed = self.preprocess.to_chw(img_processed)
+
+            processed_imgs.append(img_processed)
+            shape_lists.append(shape_list)
+
+        return processed_imgs, shape_lists