paddlex 2.0.0rc4__py3-none-any.whl → 3.0.0b2__py3-none-any.whl

paddlex/inference/pipelines/layout_parsing/layout_parsing.py

@@ -0,0 +1,362 @@
+# 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 ...results import *
+from ...components import *
+from ..ocr import OCRPipeline
+from ....utils import logging
+from ..ppchatocrv3.utils import *
+from ..table_recognition import _TableRecPipeline
+from ..table_recognition.utils import convert_4point2rect, get_ori_coordinate_for_table
+
+
+class LayoutParsingPipeline(_TableRecPipeline):
+    """Layout Analysis Pileline"""
+
+    entities = "layout_parsing"
+
+    def __init__(
+        self,
+        layout_model,
+        text_det_model,
+        text_rec_model,
+        table_model,
+        formula_rec_model,
+        doc_image_ori_cls_model=None,
+        doc_image_unwarp_model=None,
+        seal_text_det_model=None,
+        layout_batch_size=1,
+        text_det_batch_size=1,
+        text_rec_batch_size=1,
+        table_batch_size=1,
+        doc_image_ori_cls_batch_size=1,
+        doc_image_unwarp_batch_size=1,
+        seal_text_det_batch_size=1,
+        formula_rec_batch_size=1,
+        recovery=True,
+        device=None,
+        predictor_kwargs=None,
+    ):
+        super().__init__(
+            device,
+            predictor_kwargs,
+        )
+        self._build_predictor(
+            layout_model=layout_model,
+            text_det_model=text_det_model,
+            text_rec_model=text_rec_model,
+            table_model=table_model,
+            doc_image_ori_cls_model=doc_image_ori_cls_model,
+            doc_image_unwarp_model=doc_image_unwarp_model,
+            seal_text_det_model=seal_text_det_model,
+            formula_rec_model=formula_rec_model,
+        )
+        self.set_predictor(
+            layout_batch_size=layout_batch_size,
+            text_det_batch_size=text_det_batch_size,
+            text_rec_batch_size=text_rec_batch_size,
+            table_batch_size=table_batch_size,
+            doc_image_ori_cls_batch_size=doc_image_ori_cls_batch_size,
+            doc_image_unwarp_batch_size=doc_image_unwarp_batch_size,
+            seal_text_det_batch_size=seal_text_det_batch_size,
+            formula_rec_batch_size=formula_rec_batch_size,
+        )
+        self.recovery = recovery
+
+    def _build_predictor(
+        self,
+        layout_model,
+        text_det_model,
+        text_rec_model,
+        table_model,
+        formula_rec_model,
+        seal_text_det_model=None,
+        doc_image_ori_cls_model=None,
+        doc_image_unwarp_model=None,
+    ):
+        super()._build_predictor(
+            layout_model, text_det_model, text_rec_model, table_model
+        )
+
+        self.formula_predictor = self._create(formula_rec_model)
+
+        if seal_text_det_model:
+            self.curve_pipeline = self._create(
+                pipeline=OCRPipeline,
+                text_det_model=seal_text_det_model,
+                text_rec_model=text_rec_model,
+            )
+        else:
+            self.curve_pipeline = None
+        if doc_image_ori_cls_model:
+            self.oricls_predictor = self._create(doc_image_ori_cls_model)
+        else:
+            self.oricls_predictor = None
+        if doc_image_unwarp_model:
+            self.uvdoc_predictor = self._create(doc_image_unwarp_model)
+        else:
+            self.uvdoc_predictor = None
+
+        self.img_reader = ReadImage(format="BGR")
+        self.cropper = CropByBoxes()
+
+    def set_predictor(
+        self,
+        layout_batch_size=None,
+        text_det_batch_size=None,
+        text_rec_batch_size=None,
+        table_batch_size=None,
+        doc_image_ori_cls_batch_size=None,
+        doc_image_unwarp_batch_size=None,
+        seal_text_det_batch_size=None,
+        formula_rec_batch_size=None,
+        device=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
+            )
+        if table_batch_size:
+            self.table_predictor.set_predictor(batch_size=table_batch_size)
+        if formula_rec_batch_size:
+            self.formula_predictor.set_predictor(batch_size=formula_rec_batch_size)
+        if self.curve_pipeline and seal_text_det_batch_size:
+            self.curve_pipeline.text_det_model.set_predictor(
+                batch_size=seal_text_det_batch_size
+            )
+        if self.oricls_predictor and doc_image_ori_cls_batch_size:
+            self.oricls_predictor.set_predictor(batch_size=doc_image_ori_cls_batch_size)
+        if self.uvdoc_predictor and doc_image_unwarp_batch_size:
+            self.uvdoc_predictor.set_predictor(batch_size=doc_image_unwarp_batch_size)
+
+        if device:
+            if self.curve_pipeline:
+                self.curve_pipeline.set_predictor(device=device)
+            if self.oricls_predictor:
+                self.oricls_predictor.set_predictor(device=device)
+            if self.uvdoc_predictor:
+                self.uvdoc_predictor.set_predictor(device=device)
+            self.layout_predictor.set_predictor(device=device)
+            self.ocr_pipeline.set_predictor(device=device)
+
+    def predict(
+        self,
+        inputs,
+        use_doc_image_ori_cls_model=True,
+        use_doc_image_unwarp_model=True,
+        use_seal_text_det_model=True,
+        recovery=True,
+        **kwargs,
+    ):
+        self.set_predictor(**kwargs)
+        # get oricls and uvdoc results
+        img_info_list = list(self.img_reader(inputs))[0]
+        oricls_results = []
+        if self.oricls_predictor and use_doc_image_ori_cls_model:
+            oricls_results = get_oriclas_results(img_info_list, self.oricls_predictor)
+        unwarp_result = []
+        if self.uvdoc_predictor and use_doc_image_unwarp_model:
+            unwarp_result = get_unwarp_results(img_info_list, self.uvdoc_predictor)
+        img_list = [img_info["img"] for img_info in img_info_list]
+        for idx, (img_info, layout_pred) in enumerate(
+            zip(img_info_list, self.layout_predictor(img_list))
+        ):
+            page_id = idx
+            single_img_res = {
+                "input_path": "",
+                "layout_result": DetResult({}),
+                "ocr_result": OCRResult({}),
+                "table_ocr_result": [],
+                "table_result": StructureTableResult([]),
+                "layout_parsing_result": {},
+                "oricls_result": TopkResult({}),
+                "formula_result": TextRecResult({}),
+                "unwarp_result": DocTrResult({}),
+                "curve_result": [],
+            }
+            # update oricls and uvdoc result
+            if oricls_results:
+                single_img_res["oricls_result"] = oricls_results[idx]
+            if unwarp_result:
+                single_img_res["unwarp_result"] = unwarp_result[idx]
+            # update layout result
+            single_img_res["input_path"] = layout_pred["input_path"]
+            single_img_res["layout_result"] = layout_pred
+            single_img = img_info["img"]
+            table_subs = []
+            curve_subs = []
+            formula_subs = []
+            structure_res = []
+            ocr_res_with_layout = []
+            if len(layout_pred["boxes"]) > 0:
+                subs_of_img = list(self._crop_by_boxes(layout_pred))
+                # get cropped images
+                for sub in subs_of_img:
+                    box = sub["box"]
+                    xmin, ymin, xmax, ymax = [int(i) for i in box]
+                    mask_flag = True
+                    if sub["label"].lower() == "table":
+                        table_subs.append(sub)
+                    elif sub["label"].lower() == "seal":
+                        curve_subs.append(sub)
+                    elif sub["label"].lower() == "formula":
+                        formula_subs.append(sub)
+                    else:
+                        if self.recovery and recovery:
+                            # TODO: Why use the entire image?
+                            wht_im = (
+                                np.ones(single_img.shape, dtype=single_img.dtype) * 255
+                            )
+                            wht_im[ymin:ymax, xmin:xmax, :] = sub["img"]
+                            sub_ocr_res = get_ocr_res(self.ocr_pipeline, wht_im)
+                        else:
+                            sub_ocr_res = get_ocr_res(self.ocr_pipeline, sub)
+                            sub_ocr_res["dt_polys"] = get_ori_coordinate_for_table(
+                                xmin, ymin, sub_ocr_res["dt_polys"]
+                            )
+                        layout_label = sub["label"].lower()
+                        # Adapt the user label definition to specify behavior.
+                        if sub_ocr_res and sub["label"].lower() in [
+                            "image",
+                            "figure",
+                            "img",
+                            "fig",
+                        ]:
+                            get_text_in_image = kwargs.get("get_text_in_image", False)
+                            mask_flag = not get_text_in_image
+                            text_in_image = ""
+                            if get_text_in_image:
+                                text_in_image = "".join(sub_ocr_res["rec_text"])
+                                ocr_res_with_layout.append(sub_ocr_res)
+                            structure_res.append(
+                                {
+                                    "input_path": sub_ocr_res["input_path"],
+                                    "layout_bbox": box,
+                                    f"{layout_label}": {
+                                        "img": sub["img"],
+                                        f"{layout_label}_text": text_in_image,
+                                    },
+                                }
+                            )
+                        else:
+                            ocr_res_with_layout.append(sub_ocr_res)
+                            structure_res.append(
+                                {
+                                    "input_path": sub_ocr_res["input_path"],
+                                    "layout_bbox": box,
+                                    f"{layout_label}": "\n".join(
+                                        sub_ocr_res["rec_text"]
+                                    ),
+                                }
+                            )
+                    if mask_flag:
+                        single_img[ymin:ymax, xmin:xmax, :] = 255
+
+            curve_pipeline = self.ocr_pipeline
+            if self.curve_pipeline and use_seal_text_det_model:
+                curve_pipeline = self.curve_pipeline
+
+            all_curve_res = get_ocr_res(curve_pipeline, curve_subs)
+            single_img_res["curve_result"] = all_curve_res
+            if isinstance(all_curve_res, dict):
+                all_curve_res = [all_curve_res]
+            for sub, curve_res in zip(curve_subs, all_curve_res):
+                structure_res.append(
+                    {
+                        "input_path": curve_res["input_path"],
+                        "layout_bbox": sub["box"],
+                        "seal": "".join(curve_res["rec_text"]),
+                    }
+                )
+
+            all_formula_res = get_formula_res(self.formula_predictor, formula_subs)
+            single_img_res["formula_result"] = all_formula_res
+            for sub, formula_res in zip(formula_subs, all_formula_res):
+                structure_res.append(
+                    {
+                        "input_path": formula_res["input_path"],
+                        "layout_bbox": sub["box"],
+                        "formula": "".join(formula_res["rec_text"]),
+                    }
+                )
+
+            use_ocr_without_layout = kwargs.get("use_ocr_without_layout", True)
+            ocr_res = {
+                "dt_polys": [],
+                "rec_text": [],
+                "input_path": layout_pred["input_path"],
+            }
+
+            if use_ocr_without_layout:
+                ocr_res = get_ocr_res(self.ocr_pipeline, single_img)
+                ocr_res["input_path"] = layout_pred["input_path"]
+                for idx, single_dt_poly in enumerate(ocr_res["dt_polys"]):
+                    structure_res.append(
+                        {
+                            "input_path": ocr_res["input_path"],
+                            "layout_bbox": convert_4point2rect(single_dt_poly),
+                            "text_without_layout": ocr_res["rec_text"][idx],
+                        }
+                    )
+            # update ocr result
+            for layout_ocr_res in ocr_res_with_layout:
+                ocr_res["dt_polys"].extend(layout_ocr_res["dt_polys"])
+                ocr_res["rec_text"].extend(layout_ocr_res["rec_text"])
+                ocr_res["rec_score"].extend(layout_ocr_res["rec_score"])
+                ocr_res["input_path"] = single_img_res["input_path"]
+
+            all_table_ocr_res = []
+            all_table_res, _ = self.get_table_result(table_subs)
+            # get table text from html
+            structure_res_table, all_table_ocr_res = get_table_text_from_html(
+                all_table_res
+            )
+            structure_res.extend(structure_res_table)
+
+            # sort the layout result by the left top point of the box
+            structure_res = sorted_layout_boxes(structure_res, w=single_img.shape[1])
+            structure_res = LayoutParsingResult(
+                {
+                    "input_path": layout_pred["input_path"],
+                    "parsing_result": structure_res,
+                }
+            )
+
+            single_img_res["table_result"] = all_table_res
+            single_img_res["ocr_result"] = ocr_res
+            single_img_res["table_ocr_result"] = all_table_ocr_res
+            single_img_res["layout_parsing_result"] = structure_res
+            single_img_res["layout_parsing_result"]["page_id"] = page_id + 1
+            yield VisualResult(single_img_res, page_id, inputs)
+
+
+def get_formula_res(predictor, input):
+    """get formula res"""
+    res_list = []
+    if isinstance(input, list):
+        img = [im["img"] for im in input]
+    elif isinstance(input, dict):
+        img = input["img"]
+    else:
+        img = input
+    for res in predictor(img):
+        res_list.append(res)
+    return res_list

paddlex/inference/pipelines/ocr.py

@@ -0,0 +1,80 @@
+# 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 ..components import SortBoxes, CropByPolys
+from ..results import OCRResult
+from .base import BasePipeline
+from ...utils import logging
+
+
+class OCRPipeline(BasePipeline):
+    """OCR Pipeline"""
+
+    entities = "OCR"
+
+    def __init__(
+        self,
+        text_det_model,
+        text_rec_model,
+        text_det_batch_size=1,
+        text_rec_batch_size=1,
+        device=None,
+        predictor_kwargs=None,
+    ):
+        super().__init__(device, predictor_kwargs)
+        self._build_predictor(text_det_model, text_rec_model)
+        self.set_predictor(
+            text_det_batch_size=text_det_batch_size,
+            text_rec_batch_size=text_rec_batch_size,
+        )
+
+    def _build_predictor(self, text_det_model, text_rec_model):
+        self.text_det_model = self._create(model=text_det_model)
+        self.text_rec_model = self._create(model=text_rec_model)
+        self.is_curve = self.text_det_model.model_name in [
+            "PP-OCRv4_mobile_seal_det",
+            "PP-OCRv4_server_seal_det",
+        ]
+        self._sort_boxes = SortBoxes()
+        self._crop_by_polys = CropByPolys(
+            det_box_type="poly" if self.is_curve else "quad"
+        )
+
+    def set_predictor(
+        self, text_det_batch_size=None, text_rec_batch_size=None, device=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 text_rec_batch_size:
+            self.text_rec_model.set_predictor(batch_size=text_rec_batch_size)
+        if device:
+            self.text_rec_model.set_predictor(device=device)
+            self.text_det_model.set_predictor(device=device)
+
+    def predict(self, input, **kwargs):
+        self.set_predictor(**kwargs)
+        for det_res in self.text_det_model(input):
+            single_img_res = (
+                det_res if self.is_curve else next(self._sort_boxes(det_res))
+            )
+            single_img_res["rec_text"] = []
+            single_img_res["rec_score"] = []
+            if len(single_img_res["dt_polys"]) > 0:
+                all_subs_of_img = list(self._crop_by_polys(single_img_res))
+                for rec_res in self.text_rec_model(all_subs_of_img):
+                    single_img_res["rec_text"].append(rec_res["rec_text"])
+                    single_img_res["rec_score"].append(rec_res["rec_score"])
+            yield OCRResult(single_img_res)

paddlex/inference/pipelines/pp_shitu_v2.py

@@ -0,0 +1,152 @@
+# 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 pickle
+from pathlib import Path
+import numpy as np
+
+from ..utils.io import ImageReader
+from ..components import CropByBoxes, FaissIndexer
+from ..components.retrieval.faiss import FaissBuilder
+from ..results import ShiTuResult
+from .base import BasePipeline
+
+
+class ShiTuV2Pipeline(BasePipeline):
+    """ShiTuV2 Pipeline"""
+
+    entities = "PP-ShiTuV2"
+
+    def __init__(
+        self,
+        det_model,
+        rec_model,
+        det_batch_size=1,
+        rec_batch_size=1,
+        index=None,
+        score_thres=None,
+        hamming_radius=None,
+        return_k=5,
+        device=None,
+        predictor_kwargs=None,
+    ):
+        super().__init__(device, predictor_kwargs)
+        self._build_predictor(det_model, rec_model)
+        self.set_predictor(det_batch_size, rec_batch_size, device)
+        self._return_k, self._score_thres, self._hamming_radius = (
+            return_k,
+            score_thres,
+            hamming_radius,
+        )
+        self._indexer = self._build_indexer(index=index) if index else None
+
+    def _build_indexer(self, index):
+        return FaissIndexer(
+            index=index,
+            return_k=self._return_k,
+            score_thres=self._score_thres,
+            hamming_radius=self._hamming_radius,
+        )
+
+    def _build_predictor(self, det_model, rec_model):
+        self.det_model = self._create(model=det_model)
+        self.rec_model = self._create(model=rec_model)
+        self._crop_by_boxes = CropByBoxes()
+        self._img_reader = ImageReader(backend="opencv")
+
+    def set_predictor(self, det_batch_size=None, rec_batch_size=None, device=None):
+        if det_batch_size:
+            self.det_model.set_predictor(batch_size=det_batch_size)
+        if rec_batch_size:
+            self.rec_model.set_predictor(batch_size=rec_batch_size)
+        if device:
+            self.det_model.set_predictor(device=device)
+            self.rec_model.set_predictor(device=device)
+
+    def predict(self, input, index=None, **kwargs):
+        indexer = self._build_indexer(index) if index is not None else self._indexer
+        assert indexer
+        self.set_predictor(**kwargs)
+        for det_res in self.det_model(input):
+            rec_res = self.get_rec_result(det_res, indexer)
+            yield self.get_final_result(det_res, rec_res)
+
+    def get_rec_result(self, det_res, indexer):
+        if len(det_res["boxes"]) == 0:
+            full_img = self._img_reader.read(det_res["input_path"])
+            w, h = full_img.shape[:2]
+            det_res["boxes"].append(
+                {
+                    "cls_id": 0,
+                    "label": "full_img",
+                    "score": 0,
+                    "coordinate": [0, 0, h, w],
+                }
+            )
+        subs_of_img = list(self._crop_by_boxes(det_res))
+        img_list = [img["img"] for img in subs_of_img]
+        all_rec_res = list(self.rec_model(img_list))
+        all_rec_res = next(indexer(all_rec_res))
+        output = {"label": [], "score": []}
+        for res in all_rec_res:
+            output["label"].append(res["label"])
+            output["score"].append(res["score"])
+        return output
+
+    def get_final_result(self, det_res, rec_res):
+        single_img_res = {"input_path": det_res["input_path"], "boxes": []}
+        for i, obj in enumerate(det_res["boxes"]):
+            rec_scores = rec_res["score"][i]
+            labels = rec_res["label"][i]
+            single_img_res["boxes"].append(
+                {
+                    "labels": labels,
+                    "rec_scores": rec_scores,
+                    "det_score": obj["score"],
+                    "coordinate": obj["coordinate"],
+                }
+            )
+        return ShiTuResult(single_img_res)
+
+    def build_index(
+        self,
+        gallery_imgs,
+        gallery_label,
+        metric_type="IP",
+        index_type="HNSW32",
+        **kwargs
+    ):
+        return FaissBuilder.build(
+            gallery_imgs,
+            gallery_label,
+            self.rec_model.predict,
+            metric_type=metric_type,
+            index_type=index_type,
+        )
+
+    def remove_index(self, remove_ids, index):
+        return FaissBuilder.remove(remove_ids, index)
+
+    def append_index(
+        self,
+        gallery_imgs,
+        gallery_label,
+        index,
+    ):
+        return FaissBuilder.append(
+            gallery_imgs,
+            gallery_label,
+            self.rec_model.predict,
+            index,
+        )

paddlex/inference/pipelines/ppchatocrv3/__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 .ppchatocrv3 import PPChatOCRPipeline

paddlex/inference/pipelines/ppchatocrv3/ch_prompt.yaml

@@ -0,0 +1,14 @@
+kie_common_prompt:
+  task_description: '你现在的任务是从OCR文字识别的结果中提取关键词列表中每一项对应的关键信息。
+      OCR的文字识别结果使用```符号包围，包含所识别出来的文字，顺序在原始图片中从左至右、从上至下。
+      我指定的关键词列表使用[]符号包围。请注意OCR的文字识别结果可能存在长句子换行被切断、不合理的分词、
+      文字被错误合并等问题，你需要结合上下文语义进行综合判断，以抽取准确的关键信息。'
+  output_format: '在返回结果时使用JSON格式，包含多个key-value对，key值为我指定的关键词，value值为所抽取的结果。
+      如果认为OCR识别结果中没有关键词key对应的value，则将value赋值为"未知"。请只输出json格式的结果，
+      并做json格式校验后返回，不要包含其它多余文字！'
+kie_table_prompt:
+  task_description: '你现在的任务是从输入的html格式的表格内容中提取关键词列表中每一项对应的关键信息,
+      表格内容用```符号包围，我指定的关键词列表使用[]符号包围。你需要结合上下文语义进行综合判断，以抽取准确的关键信息。
+      在返回结果时使用JSON格式，包含多个key-value对，key值为我指定的关键词，value值为所抽取的结果。
+      如果认为输入的表格内容中没有关键词key对应的value值，则将value赋值为"未知"。
+      请只输出json格式的结果，并做json格式校验后返回，不要包含其它多余文字！'