Sunting78 1 жил өмнө
parent
commit
622cc3091e

+ 1 - 0
paddlex/inference/pipelines/__init__.py

@@ -33,6 +33,7 @@ from .single_model_pipeline import (
 )
 from .ocr import OCRPipeline
 from .table_recognition import TableRecPipeline
+from .seal_recognition import SealTextRecPipeline
 from .ppchatocrv3 import PPChatOCRPipeline
 
 

+ 15 - 0
paddlex/inference/pipelines/seal_recognition/__init__.py

@@ -0,0 +1,15 @@
+# copyright (c) 2024 PaddlePaddle Authors. All Rights Reserve.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .seal_recognition import SealTextRecPipeline

+ 104 - 0
paddlex/inference/pipelines/seal_recognition/seal_recognition.py

@@ -0,0 +1,104 @@
+# copyright (c) 2024 PaddlePaddle Authors. All Rights Reserve.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import numpy as np
+from .utils import *
+from ..base import BasePipeline
+from ..ocr import OCRPipeline
+from ....utils import logging
+from ...components import CropByBoxes
+from ...results import SealResult
+
+
+class SealTextRecPipeline(BasePipeline):
+    """Seal Recognition Pipeline"""
+
+    entities = "seal_recognition"
+
+    def __init__(
+        self,
+        layout_model,
+        text_det_model,
+        text_rec_model,
+        layout_batch_size=1,
+        text_det_batch_size=1,
+        text_rec_batch_size=1,
+        predictor_kwargs=None,
+    ):
+        self.layout_model = layout_model
+        self.text_det_model = text_det_model
+        self.text_rec_model = text_rec_model
+        self.layout_batch_size = layout_batch_size
+        self.text_det_batch_size = text_det_batch_size
+        self.text_rec_batch_size = text_rec_batch_size
+        self.predictor_kwargs = predictor_kwargs
+        super().__init__(predictor_kwargs=predictor_kwargs)
+        self._build_predictor()
+
+    def _build_predictor(
+        self,
+    ):
+        self.layout_predictor = self._create_model(model=self.layout_model)
+        self.ocr_pipeline = OCRPipeline(
+            text_det_model=self.text_det_model,
+            text_rec_model=self.text_rec_model,
+            text_det_batch_size=self.text_det_batch_size,
+            text_rec_batch_size=self.text_rec_batch_size,
+            predictor_kwargs=self.predictor_kwargs,
+        )
+        self._crop_by_boxes = CropByBoxes()
+        self.layout_predictor.set_predictor(batch_size=self.layout_batch_size)
+        self.ocr_pipeline.text_rec_model.set_predictor(
+            batch_size=self.text_rec_batch_size
+        )
+
+    def set_predictor(
+        self,
+        layout_batch_size=None,
+        text_det_batch_size=None,
+        text_rec_batch_size=None,
+    ):
+        if text_det_batch_size and text_det_batch_size > 1:
+            logging.warning(
+                f"text det model only support batch_size=1 now,the setting of text_det_batch_size={text_det_batch_size} will not using! "
+            )
+        if layout_batch_size:
+            self.layout_predictor.set_predictor(batch_size=layout_batch_size)
+        if text_rec_batch_size:
+            self.ocr_pipeline.text_rec_model.set_predictor(
+                batch_size=text_rec_batch_size
+            )
+
+    def predict(self, x):
+        for layout_pred in self.layout_predictor(x):
+            single_img_res = {
+                "input_path": "",
+                "layout_result": {},
+                "ocr_result": {},
+            }
+            # update layout result
+            single_img_res["input_path"] = layout_pred["input_path"]
+            single_img_res["layout_result"] = layout_pred
+
+            seal_subs = []
+            if len(layout_pred["boxes"]) > 0:
+                subs_of_img = list(self._crop_by_boxes(layout_pred))
+                # get cropped images with label "seal"
+                for sub in subs_of_img:
+                    box = sub["box"]
+                    if sub["label"].lower() == "seal":
+                        seal_subs.append(sub)
+            all_seal_ocr_res = get_ocr_res(self.ocr_pipeline, seal_subs)
+            single_img_res["ocr_result"] = all_seal_ocr_res
+            yield SealResult(single_img_res)

+ 16 - 0
paddlex/inference/pipelines/seal_recognition/utils.py

@@ -0,0 +1,16 @@
+import os
+from pathlib import Path
+
+
+def get_ocr_res(pipeline, input):
+    """get ocr res"""
+    ocr_res_list = []
+    if isinstance(input, list):
+        img = [im["img"] for im in input]
+    elif isinstance(input, dict):
+        img = input["img"]
+    else:
+        img = input
+    for ocr_res in pipeline(img):
+        ocr_res_list.append(ocr_res)
+    return ocr_res_list

+ 1 - 0
paddlex/inference/results/__init__.py

@@ -17,6 +17,7 @@ from .clas import TopkResult, MLClassResult
 from .text_det import TextDetResult
 from .text_rec import TextRecResult
 from .table_rec import TableRecResult, StructureTableResult, TableResult
+from .seal_rec import SealResult
 from .ocr import OCRResult
 from .det import DetResult
 from .seg import SegResult

+ 19 - 0
paddlex/inference/results/seal_rec.py

@@ -0,0 +1,19 @@
+from pathlib import Path
+from .base import BaseResult
+
+
+class SealResult(BaseResult):
+    """SealResult"""
+
+    def save_to_img(self, save_path):
+        if not save_path.lower().endswith((".jpg", ".png")):
+            input_path = self["input_path"]
+            save_path = Path(save_path) / f"{Path(input_path).stem}"
+        else:
+            save_path = Path(save_path).stem
+        layout_save_path = f"{save_path}_layout.jpg"
+        layout_result = self["layout_result"]
+        layout_result.save_to_img(layout_save_path)
+        for idx, seal_result in enumerate(self["ocr_result"]):
+            ocr_save_path = f"{save_path}_{idx}_seal_ocr.jpg"
+            seal_result.save_to_img(ocr_save_path)

+ 10 - 0
paddlex/pipelines/seal_recognition.yaml

@@ -0,0 +1,10 @@
+Global:
+  pipeline_name: seal_recognition
+  input: https://paddle-model-ecology.bj.bcebos.com/paddlex/imgs/demo_image/seal_text_det.png
+
+Pipeline:
+  layout_model: RT-DETR-H_layout_3cls
+  text_det_model: PP-OCRv4_server_seal_det
+  text_rec_model: PP-OCRv4_server_rec
+  layout_batch_size: 1
+  text_rec_batch_size: 1