paddlex 3.0.0rc1__py3-none-any.whl → 3.0.2__py3-none-any.whl

paddlex/inference/pipelines/seal_recognition/pipeline.py

@@ -23,17 +23,15 @@ from ...common.reader import ReadImage
 from ...models.object_detection.result import DetResult
 from ...utils.hpi import HPIConfig
 from ...utils.pp_option import PaddlePredictorOption
+from .._parallel import AutoParallelImageSimpleInferencePipeline
 from ..base import BasePipeline
 from ..components import CropByBoxes
 from .result import SealRecognitionResult
 
 
-@pipeline_requires_extra("ocr")
-class SealRecognitionPipeline(BasePipeline):
+class _SealRecognitionPipeline(BasePipeline):
     """Seal Recognition Pipeline"""
 
-    entities = ["seal_recognition"]
-
     def __init__(
         self,
         config: Dict,
@@ -49,9 +47,9 @@ class SealRecognitionPipeline(BasePipeline):
             device (str, optional): Device to run the predictions on. Defaults to None.
             pp_option (PaddlePredictorOption, optional): PaddlePredictor options. Defaults to None.
             use_hpip (bool, optional): Whether to use the high-performance
-                inference plugin (HPIP). Defaults to False.
+                inference plugin (HPIP) by default. Defaults to False.
             hpi_config (Optional[Union[Dict[str, Any], HPIConfig]], optional):
-                The high-performance inference configuration dictionary.
+                The default high-performance inference configuration dictionary.
                 Defaults to None.
         """
 
@@ -104,7 +102,7 @@ class SealRecognitionPipeline(BasePipeline):
 
         self._crop_by_boxes = CropByBoxes()
 
-        self.batch_sampler = ImageBatchSampler(batch_size=1)
+        self.batch_sampler = ImageBatchSampler(batch_size=config.get("batch_size", 1))
 
         self.img_reader = ReadImage(format="BGR")
 
@@ -180,7 +178,7 @@ class SealRecognitionPipeline(BasePipeline):
         use_doc_orientation_classify: Optional[bool] = None,
         use_doc_unwarping: Optional[bool] = None,
         use_layout_detection: Optional[bool] = None,
-        layout_det_res: Optional[DetResult] = None,
+        layout_det_res: Optional[Union[DetResult, List[DetResult]]] = None,
         layout_threshold: Optional[Union[float, dict]] = None,
         layout_nms: Optional[bool] = None,
         layout_unclip_ratio: Optional[Union[float, Tuple[float, float]]] = None,
@@ -201,29 +199,38 @@ class SealRecognitionPipeline(BasePipeline):
         if not self.check_model_settings_valid(model_settings, layout_det_res):
             yield {"error": "the input params for model settings are invalid!"}
 
-        for img_id, batch_data in enumerate(self.batch_sampler(input)):
-            image_array = self.img_reader(batch_data.instances)[0]
+        external_layout_det_results = layout_det_res
+        if external_layout_det_results is not None:
+            if not isinstance(external_layout_det_results, list):
+                external_layout_det_results = [external_layout_det_results]
+            external_layout_det_results = iter(external_layout_det_results)
+
+        for _, batch_data in enumerate(self.batch_sampler(input)):
+            image_arrays = self.img_reader(batch_data.instances)
 
             if model_settings["use_doc_preprocessor"]:
-                doc_preprocessor_res = next(
+                doc_preprocessor_results = list(
                     self.doc_preprocessor_pipeline(
-                        image_array,
+                        image_arrays,
                         use_doc_orientation_classify=use_doc_orientation_classify,
                         use_doc_unwarping=use_doc_unwarping,
                     )
                 )
             else:
-                doc_preprocessor_res = {"output_img": image_array}
+                doc_preprocessor_results = [{"output_img": arr} for arr in image_arrays]
 
-            doc_preprocessor_image = doc_preprocessor_res["output_img"]
+            doc_preprocessor_images = [
+                item["output_img"] for item in doc_preprocessor_results
+            ]
 
-            seal_res_list = []
-            seal_region_id = 1
-            if not model_settings["use_layout_detection"] and layout_det_res is None:
-                layout_det_res = {}
-                seal_ocr_res = next(
+            if (
+                not model_settings["use_layout_detection"]
+                and external_layout_det_results is None
+            ):
+                layout_det_results = [{} for _ in doc_preprocessor_images]
+                flat_seal_results = list(
                     self.seal_ocr_pipeline(
-                        doc_preprocessor_image,
+                        doc_preprocessor_images,
                         text_det_limit_side_len=seal_det_limit_side_len,
                         text_det_limit_type=seal_det_limit_type,
                         text_det_thresh=seal_det_thresh,
@@ -232,48 +239,97 @@ class SealRecognitionPipeline(BasePipeline):
                         text_rec_score_thresh=seal_rec_score_thresh,
                     )
                 )
-                seal_ocr_res["seal_region_id"] = seal_region_id
-                seal_res_list.append(seal_ocr_res)
-                seal_region_id += 1
+                for seal_res in flat_seal_results:
+                    seal_res["seal_region_id"] = 1
+                seal_results = [[item] for item in flat_seal_results]
             else:
                 if model_settings["use_layout_detection"]:
-                    layout_det_res = next(
+                    layout_det_results = list(
                         self.layout_det_model(
-                            doc_preprocessor_image,
+                            doc_preprocessor_images,
                             threshold=layout_threshold,
                             layout_nms=layout_nms,
                             layout_unclip_ratio=layout_unclip_ratio,
                             layout_merge_bboxes_mode=layout_merge_bboxes_mode,
                         )
                     )
-
-                for box_info in layout_det_res["boxes"]:
-                    if box_info["label"].lower() in ["seal"]:
-                        crop_img_info = self._crop_by_boxes(
-                            doc_preprocessor_image, [box_info]
-                        )
-                        crop_img_info = crop_img_info[0]
-                        seal_ocr_res = next(
-                            self.seal_ocr_pipeline(
-                                crop_img_info["img"],
-                                text_det_limit_side_len=seal_det_limit_side_len,
-                                text_det_limit_type=seal_det_limit_type,
-                                text_det_thresh=seal_det_thresh,
-                                text_det_box_thresh=seal_det_box_thresh,
-                                text_det_unclip_ratio=seal_det_unclip_ratio,
-                                text_rec_score_thresh=seal_rec_score_thresh,
+                else:
+                    layout_det_results = []
+                    for _ in doc_preprocessor_images:
+                        try:
+                            layout_det_res = next(external_layout_det_results)
+                        except StopIteration:
+                            raise ValueError("No more layout det results")
+                        layout_det_results.append(layout_det_res)
+
+                cropped_imgs = []
+                chunk_indices = [0]
+                for doc_preprocessor_image, layout_det_res in zip(
+                    doc_preprocessor_images, layout_det_results
+                ):
+                    for box_info in layout_det_res["boxes"]:
+                        if box_info["label"].lower() in ["seal"]:
+                            crop_img_info = self._crop_by_boxes(
+                                doc_preprocessor_image, [box_info]
                             )
-                        )
-                        seal_ocr_res["seal_region_id"] = seal_region_id
-                        seal_res_list.append(seal_ocr_res)
+                            crop_img_info = crop_img_info[0]
+                            cropped_imgs.append(crop_img_info["img"])
+                    chunk_indices.append(len(cropped_imgs))
+
+                flat_seal_results = list(
+                    self.seal_ocr_pipeline(
+                        cropped_imgs,
+                        text_det_limit_side_len=seal_det_limit_side_len,
+                        text_det_limit_type=seal_det_limit_type,
+                        text_det_thresh=seal_det_thresh,
+                        text_det_box_thresh=seal_det_box_thresh,
+                        text_det_unclip_ratio=seal_det_unclip_ratio,
+                        text_rec_score_thresh=seal_rec_score_thresh,
+                    )
+                )
+
+                seal_results = [
+                    flat_seal_results[i:j]
+                    for i, j in zip(chunk_indices[:-1], chunk_indices[1:])
+                ]
+
+                for seal_results_for_img in seal_results:
+                    seal_region_id = 1
+                    for seal_res in seal_results_for_img:
+                        seal_res["seal_region_id"] = seal_region_id
                         seal_region_id += 1
 
-            single_img_res = {
-                "input_path": batch_data.input_paths[0],
-                "page_index": batch_data.page_indexes[0],
-                "doc_preprocessor_res": doc_preprocessor_res,
-                "layout_det_res": layout_det_res,
-                "seal_res_list": seal_res_list,
-                "model_settings": model_settings,
-            }
-            yield SealRecognitionResult(single_img_res)
+            for (
+                input_path,
+                page_index,
+                doc_preprocessor_res,
+                layout_det_res,
+                seal_results_for_img,
+            ) in zip(
+                batch_data.input_paths,
+                batch_data.page_indexes,
+                doc_preprocessor_results,
+                layout_det_results,
+                seal_results,
+            ):
+                single_img_res = {
+                    "input_path": input_path,
+                    "page_index": page_index,
+                    "doc_preprocessor_res": doc_preprocessor_res,
+                    "layout_det_res": layout_det_res,
+                    "seal_res_list": seal_results_for_img,
+                    "model_settings": model_settings,
+                }
+                yield SealRecognitionResult(single_img_res)
+
+
+@pipeline_requires_extra("ocr")
+class SealRecognitionPipeline(AutoParallelImageSimpleInferencePipeline):
+    entities = ["seal_recognition"]
+
+    @property
+    def _pipeline_cls(self):
+        return _SealRecognitionPipeline
+
+    def _get_batch_size(self, config):
+        return config.get("batch_size", 1)

paddlex/inference/pipelines/semantic_segmentation/pipeline.py

@@ -20,15 +20,13 @@ from ....utils.deps import pipeline_requires_extra
 from ...models.semantic_segmentation.result import SegResult
 from ...utils.hpi import HPIConfig
 from ...utils.pp_option import PaddlePredictorOption
+from .._parallel import AutoParallelImageSimpleInferencePipeline
 from ..base import BasePipeline
 
 
-@pipeline_requires_extra("cv")
-class SemanticSegmentationPipeline(BasePipeline):
+class _SemanticSegmentationPipeline(BasePipeline):
     """Semantic Segmentation Pipeline"""
 
-    entities = "semantic_segmentation"
-
     def __init__(
         self,
         config: Dict,
@@ -45,9 +43,9 @@ class SemanticSegmentationPipeline(BasePipeline):
             device (str): The device to run the prediction on. Default is None.
             pp_option (PaddlePredictorOption): Options for PaddlePaddle predictor. Default is None.
             use_hpip (bool, optional): Whether to use the high-performance
-                inference plugin (HPIP). Defaults to False.
+                inference plugin (HPIP) by default. Defaults to False.
             hpi_config (Optional[Union[Dict[str, Any], HPIConfig]], optional):
-                The high-performance inference configuration dictionary.
+                The default high-performance inference configuration dictionary.
                 Defaults to None.
         """
         super().__init__(
@@ -83,3 +81,15 @@ class SemanticSegmentationPipeline(BasePipeline):
             SegResult: The predicted segmentation results.
         """
         yield from self.semantic_segmentation_model(input, target_size=target_size)
+
+
+@pipeline_requires_extra("cv")
+class SemanticSegmentationPipeline(AutoParallelImageSimpleInferencePipeline):
+    entities = "semantic_segmentation"
+
+    @property
+    def _pipeline_cls(self):
+        return _SemanticSegmentationPipeline
+
+    def _get_batch_size(self, config):
+        return config["SubModules"]["SemanticSegmentation"].get("batch_size", 1)

paddlex/inference/pipelines/small_object_detection/pipeline.py

@@ -20,15 +20,13 @@ from ....utils.deps import pipeline_requires_extra
 from ...models.object_detection.result import DetResult
 from ...utils.hpi import HPIConfig
 from ...utils.pp_option import PaddlePredictorOption
+from .._parallel import AutoParallelImageSimpleInferencePipeline
 from ..base import BasePipeline
 
 
-@pipeline_requires_extra("cv")
-class SmallObjectDetectionPipeline(BasePipeline):
+class _SmallObjectDetectionPipeline(BasePipeline):
     """Small Object Detection Pipeline"""
 
-    entities = "small_object_detection"
-
     def __init__(
         self,
         config: Dict,
@@ -45,9 +43,9 @@ class SmallObjectDetectionPipeline(BasePipeline):
             device (str): The device to run the prediction on. Default is None.
             pp_option (PaddlePredictorOption): Options for PaddlePaddle predictor. Default is None.
             use_hpip (bool, optional): Whether to use the high-performance
-                inference plugin (HPIP). Defaults to False.
+                inference plugin (HPIP) by default. Defaults to False.
             hpi_config (Optional[Union[Dict[str, Any], HPIConfig]], optional):
-                The high-performance inference configuration dictionary.
+                The default high-performance inference configuration dictionary.
                 Defaults to None.
         """
         super().__init__(
@@ -83,3 +81,15 @@ class SmallObjectDetectionPipeline(BasePipeline):
             DetResult: The predicted small object detection results.
         """
         yield from self.small_object_detection_model(input, threshold=threshold)
+
+
+@pipeline_requires_extra("cv")
+class SmallObjectDetectionPipeline(AutoParallelImageSimpleInferencePipeline):
+    entities = "small_object_detection"
+
+    @property
+    def _pipeline_cls(self):
+        return _SmallObjectDetectionPipeline
+
+    def _get_batch_size(self, config):
+        return config["SubModules"]["SmallObjectDetection"].get("batch_size", 1)

paddlex/inference/pipelines/table_recognition/pipeline.py

@@ -24,6 +24,7 @@ from ...common.reader import ReadImage
 from ...models.object_detection.result import DetResult
 from ...utils.hpi import HPIConfig
 from ...utils.pp_option import PaddlePredictorOption
+from .._parallel import AutoParallelImageSimpleInferencePipeline
 from ..base import BasePipeline
 from ..components import CropByBoxes
 from ..doc_preprocessor.result import DocPreprocessorResult
@@ -33,12 +34,9 @@ from .table_recognition_post_processing import get_table_recognition_res
 from .utils import get_neighbor_boxes_idx
 
 
-@pipeline_requires_extra("ocr")
-class TableRecognitionPipeline(BasePipeline):
+class _TableRecognitionPipeline(BasePipeline):
     """Table Recognition Pipeline"""
 
-    entities = ["table_recognition"]
-
     def __init__(
         self,
         config: Dict,
@@ -54,9 +52,9 @@ class TableRecognitionPipeline(BasePipeline):
             device (str, optional): Device to run the predictions on. Defaults to None.
             pp_option (PaddlePredictorOption, optional): PaddlePredictor options. Defaults to None.
             use_hpip (bool, optional): Whether to use the high-performance
-                inference plugin (HPIP). Defaults to False.
+                inference plugin (HPIP) by default. Defaults to False.
             hpi_config (Optional[Union[Dict[str, Any], HPIConfig]], optional):
-                The high-performance inference configuration dictionary.
+                The default high-performance inference configuration dictionary.
                 Defaults to None.
         """
 
@@ -290,7 +288,7 @@ class TableRecognitionPipeline(BasePipeline):
         image_array: np.ndarray,
         overall_ocr_res: OCRResult,
         table_box: list,
-        use_table_cells_ocr_results: bool = False,
+        use_ocr_results_with_table_cells: bool = False,
         flag_find_nei_text: bool = True,
         cell_sort_by_y_projection: bool = False,
     ) -> SingleTableRecognitionResult:
@@ -302,17 +300,15 @@ class TableRecognitionPipeline(BasePipeline):
             overall_ocr_res (OCRResult): Overall OCR result obtained after running the OCR pipeline.
                 The overall OCR results containing text recognition information.
             table_box (list): The table box coordinates.
-            use_table_cells_ocr_results (bool): whether to use OCR results with cells.
+            use_ocr_results_with_table_cells (bool): whether to use OCR results with cells.
             flag_find_nei_text (bool): Whether to find neighboring text.
             cell_sort_by_y_projection (bool): Whether to sort the matched OCR boxes by y-projection.
         Returns:
             SingleTableRecognitionResult: single table recognition result.
         """
         table_structure_pred = next(self.table_structure_model(image_array))
-        if use_table_cells_ocr_results == True:
-            table_cells_result = list(
-                map(lambda arr: arr.tolist(), table_structure_pred["bbox"])
-            )
+        if use_ocr_results_with_table_cells == True:
+            table_cells_result = table_structure_pred["bbox"]
             table_cells_result = [
                 [rect[0], rect[1], rect[4], rect[5]] for rect in table_cells_result
             ]
@@ -326,7 +322,7 @@ class TableRecognitionPipeline(BasePipeline):
             table_structure_pred,
             overall_ocr_res,
             cells_texts_list,
-            use_table_cells_ocr_results,
+            use_ocr_results_with_table_cells,
             cell_sort_by_y_projection=cell_sort_by_y_projection,
         )
         neighbor_text = ""
@@ -355,7 +351,7 @@ class TableRecognitionPipeline(BasePipeline):
         text_det_box_thresh: Optional[float] = None,
         text_det_unclip_ratio: Optional[float] = None,
         text_rec_score_thresh: Optional[float] = None,
-        use_table_cells_ocr_results: bool = False,
+        use_ocr_results_with_table_cells: bool = False,
         cell_sort_by_y_projection: Optional[bool] = None,
         **kwargs,
     ) -> TableRecognitionResult:
@@ -371,7 +367,7 @@ class TableRecognitionPipeline(BasePipeline):
                 It will be used if it is not None and use_ocr_model is False.
             layout_det_res (DetResult): The layout detection result.
                 It will be used if it is not None and use_layout_detection is False.
-            use_table_cells_ocr_results (bool): whether to use OCR results with cells.
+            use_ocr_results_with_table_cells (bool): whether to use OCR results with cells.
             cell_sort_by_y_projection (bool): Whether to sort the matched OCR boxes by y-projection.
             **kwargs: Additional keyword arguments.
 
@@ -421,7 +417,7 @@ class TableRecognitionPipeline(BasePipeline):
                         text_rec_score_thresh=text_rec_score_thresh,
                     )
                 )
-            elif use_table_cells_ocr_results == True:
+            elif use_ocr_results_with_table_cells == True:
                 assert self.general_ocr_config_bak != None
                 self.general_ocr_pipeline = self.create_pipeline(
                     self.general_ocr_config_bak
@@ -437,7 +433,7 @@ class TableRecognitionPipeline(BasePipeline):
                     doc_preprocessor_image,
                     overall_ocr_res,
                     table_box,
-                    use_table_cells_ocr_results,
+                    use_ocr_results_with_table_cells,
                     flag_find_nei_text=False,
                     cell_sort_by_y_projection=cell_sort_by_y_projection,
                 )
@@ -458,7 +454,7 @@ class TableRecognitionPipeline(BasePipeline):
                                 crop_img_info["img"],
                                 overall_ocr_res,
                                 table_box,
-                                use_table_cells_ocr_results,
+                                use_ocr_results_with_table_cells,
                                 cell_sort_by_y_projection=cell_sort_by_y_projection,
                             )
                         )
@@ -476,3 +472,15 @@ class TableRecognitionPipeline(BasePipeline):
                 "model_settings": model_settings,
             }
             yield TableRecognitionResult(single_img_res)
+
+
+@pipeline_requires_extra("ocr")
+class TableRecognitionPipeline(AutoParallelImageSimpleInferencePipeline):
+    entities = ["table_recognition"]
+
+    @property
+    def _pipeline_cls(self):
+        return _TableRecognitionPipeline
+
+    def _get_batch_size(self, config):
+        return 1