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

paddlex/inference/pipelines/serving/models.py

@@ -0,0 +1,30 @@
+# 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 typing import Generic, TypeVar
+
+from pydantic import BaseModel
+
+
+class Response(BaseModel):
+    logId: str
+    errorCode: int
+    errorMsg: str
+
+
+ResultT = TypeVar("ResultT", bound=BaseModel)
+
+
+class ResultResponse(Response, Generic[ResultT]):
+    result: ResultT

paddlex/inference/pipelines/serving/server.py

@@ -0,0 +1,25 @@
+# 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 uvicorn
+from fastapi import FastAPI
+
+
+def run_server(app: FastAPI, *, host: str, port: int, debug: bool) -> None:
+    # XXX: Currently, `debug` is not used.
+    # HACK: Fix duplicate logs
+    uvicorn_version = tuple(int(x) for x in uvicorn.__version__.split("."))
+    if uvicorn_version < (0, 19, 0):
+        uvicorn.config.LOGGING_CONFIG["loggers"]["uvicorn"]["propagate"] = False
+    uvicorn.run(app, host=host, port=port, log_level="info")

paddlex/inference/pipelines/serving/storage.py

@@ -0,0 +1,161 @@
+# 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 abc
+from os import PathLike
+from pathlib import Path
+from typing import Any, Dict, Literal, Optional, Union, Protocol, runtime_checkable
+
+from baidubce.auth.bce_credentials import BceCredentials
+from baidubce.bce_client_configuration import BceClientConfiguration
+from baidubce.services.bos.bos_client import BosClient
+from pydantic import BaseModel, Discriminator, SecretStr, TypeAdapter
+from typing_extensions import Annotated, assert_never
+
+
+class InMemoryStorageConfig(BaseModel):
+    type: Literal["memory"] = "memory"
+
+
+class FileSystemStorageConfig(BaseModel):
+    directory: Union[str, PathLike]
+
+    type: Literal["file_system"] = "file_system"
+
+
+class BOSConfig(BaseModel):
+    endpoint: str
+    ak: SecretStr
+    sk: SecretStr
+    bucket_name: str
+    key_prefix: Optional[str] = None
+    connection_timeout_in_mills: Optional[int] = None
+
+    type: Literal["bos"] = "bos"
+
+
+FileStorageConfig = Union[InMemoryStorageConfig, FileSystemStorageConfig, BOSConfig]
+
+
+@runtime_checkable
+class SupportsGetURL(Protocol):
+    def get_url(self, key: str) -> str: ...
+
+
+class Storage(metaclass=abc.ABCMeta):
+    @abc.abstractmethod
+    def get(self, key: str) -> bytes:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def set(self, key: str, value: bytes) -> None:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def delete(self, key: str) -> None:
+        raise NotImplementedError
+
+
+class InMemoryStorage(Storage):
+    def __init__(self, config: InMemoryStorageConfig) -> None:
+        super().__init__()
+        self._data: Dict[str, bytes] = {}
+
+    def get(self, key: str) -> bytes:
+        return self._data[key]
+
+    def set(self, key: str, value: bytes) -> None:
+        self._data[key] = value
+
+    def delete(self, key: str) -> None:
+        del self._data[key]
+
+
+class FileSystemStorage(Storage):
+    def __init__(self, config: FileSystemStorageConfig) -> None:
+        super().__init__()
+        self._directory = Path(config.directory)
+        self._directory.mkdir(exist_ok=True)
+
+    def get(self, key: str) -> bytes:
+        with open(self._get_file_path(key), "rb") as f:
+            contents = f.read()
+        return contents
+
+    def set(self, key: str, value: bytes) -> None:
+        path = self._get_file_path(key)
+        path.parent.mkdir(exist_ok=True)
+        with open(path, "wb") as f:
+            f.write(value)
+
+    def delete(self, key: str) -> None:
+        file_path = self._get_file_path(key)
+        file_path.unlink(missing_ok=True)
+
+    def _get_file_path(self, key: str) -> Path:
+        return self._directory / key
+
+
+class BOS(Storage):
+    def __init__(self, config: BOSConfig) -> None:
+        super().__init__()
+        bos_cfg = BceClientConfiguration(
+            credentials=BceCredentials(
+                config.ak.get_secret_value(), config.sk.get_secret_value()
+            ),
+            endpoint=config.endpoint,
+            connection_timeout_in_mills=config.connection_timeout_in_mills,
+        )
+        self._client = BosClient(bos_cfg)
+        self._bucket_name = config.bucket_name
+        self._key_prefix = config.key_prefix
+
+    def get(self, key: str) -> bytes:
+        key = self._get_full_key(key)
+        return self._client.get_object_as_string(bucket_name=self._bucket_name, key=key)
+
+    def set(self, key: str, value: bytes) -> None:
+        key = self._get_full_key(key)
+        self._client.put_object_from_string(
+            bucket=self._bucket_name, key=key, data=value
+        )
+
+    def delete(self, key: str) -> None:
+        key = self._get_full_key(key)
+        self._client.delete_object(bucket_name=self._bucket_name, key=key)
+
+    def get_url(self, key: str) -> str:
+        key = self._get_full_key(key)
+        return self._client.generate_pre_signed_url(
+            self._bucket_name, key, expiration_in_seconds=-1
+        ).decode("ascii")
+
+    def _get_full_key(self, key: str) -> str:
+        if self._key_prefix:
+            return f"{self._key_prefix}/{key}"
+        return key
+
+
+def create_storage(dic: Dict[str, Any], /) -> Storage:
+    config = TypeAdapter(
+        Annotated[FileStorageConfig, Discriminator("type")]
+    ).validate_python(dic)
+    if config.type == "memory":
+        return InMemoryStorage(config)
+    elif config.type == "file_system":
+        return FileSystemStorage(config)
+    elif config.type == "bos":
+        return BOS(config)
+    else:
+        assert_never(config)

paddlex/inference/pipelines/serving/utils.py

@@ -0,0 +1,190 @@
+# 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 asyncio
+import base64
+import io
+import os
+import re
+import uuid
+from functools import partial
+from typing import Awaitable, Callable, List, Literal, Optional, TypeVar, Final, Tuple
+from urllib.parse import parse_qs, urlparse
+
+import aiohttp
+import cv2
+import fitz
+import numpy as np
+import pandas as pd
+import yarl
+from PIL import Image
+from typing_extensions import ParamSpec, assert_never
+
+FileType = Literal["IMAGE", "PDF"]
+
+_P = ParamSpec("_P")
+_R = TypeVar("_R")
+
+
+def generate_log_id() -> str:
+    return str(uuid.uuid4())
+
+
+def generate_request_id() -> str:
+    return str(uuid.uuid4())
+
+
+def is_url(s: str) -> bool:
+    if not (s.startswith("http://") or s.startswith("https://")):
+        # Quick rejection
+        return False
+    result = urlparse(s)
+    return all([result.scheme, result.netloc]) and result.scheme in ("http", "https")
+
+
+def infer_file_type(url: str) -> FileType:
+    # Is it more reliable to guess the file type based on the response headers?
+    SUPPORTED_IMG_EXTS: Final[List[str]] = [".jpg", ".jpeg", ".png"]
+
+    url_parts = urlparse(url)
+    ext = os.path.splitext(url_parts.path)[1]
+    # HACK: The support for BOS URLs with query params is implementation-based,
+    # not interface-based.
+    is_bos_url = (
+        re.fullmatch(r"(?:bj|bd|su|gz|cd|hkg|fwh|fsh)\.bcebos\.com", url_parts.netloc)
+        is not None
+    )
+    if is_bos_url and url_parts.query:
+        params = parse_qs(url_parts.query)
+        if (
+            "responseContentDisposition" not in params
+            or len(params["responseContentDisposition"]) != 1
+        ):
+            raise ValueError("`responseContentDisposition` not found")
+        match_ = re.match(
+            r"attachment;filename=(.*)", params["responseContentDisposition"][0]
+        )
+        if not match_ or not match_.groups()[0] is not None:
+            raise ValueError(
+                "Failed to extract the filename from `responseContentDisposition`"
+            )
+        ext = os.path.splitext(match_.groups()[0])[1]
+    ext = ext.lower()
+    if ext == ".pdf":
+        return "PDF"
+    elif ext in SUPPORTED_IMG_EXTS:
+        return "IMAGE"
+    else:
+        raise ValueError("Unsupported file type")
+
+
+async def get_raw_bytes(file: str, session: aiohttp.ClientSession) -> bytes:
+    if is_url(file):
+        async with session.get(yarl.URL(file, encoded=True)) as resp:
+            return await resp.read()
+    else:
+        return base64.b64decode(file)
+
+
+def image_bytes_to_array(data: bytes) -> np.ndarray:
+    return cv2.imdecode(np.frombuffer(data, np.uint8), cv2.IMREAD_COLOR)
+
+
+def image_bytes_to_image(data: bytes) -> Image.Image:
+    return Image.open(io.BytesIO(data))
+
+
+def image_to_bytes(image: Image.Image, format: str = "JPEG") -> bytes:
+    with io.BytesIO() as f:
+        image.save(f, format=format)
+        img_bytes = f.getvalue()
+    return img_bytes
+
+
+def image_array_to_bytes(image: np.ndarray, ext: str = ".jpg") -> str:
+    image = cv2.imencode(ext, image)[1]
+    return image.tobytes()
+
+
+def csv_bytes_to_data_frame(data: bytes) -> pd.DataFrame:
+    with io.StringIO(data.decode("utf-8")) as f:
+        df = pd.read_csv(f)
+    return df
+
+
+def data_frame_to_bytes(df: str) -> str:
+    return df.to_csv().encode("utf-8")
+
+
+def base64_encode(data: bytes) -> str:
+    return base64.b64encode(data).decode("ascii")
+
+
+def read_pdf(
+    bytes_: bytes, resize: bool = False, max_num_imgs: Optional[int] = None
+) -> List[np.ndarray]:
+    images: List[np.ndarray] = []
+    img_size = None
+    with fitz.open("pdf", bytes_) as doc:
+        for page in doc:
+            if max_num_imgs is not None and len(images) >= max_num_imgs:
+                break
+            # TODO: Do not always use zoom=2.0
+            zoom = 2.0
+            deg = 0
+            mat = fitz.Matrix(zoom, zoom).prerotate(deg)
+            pixmap = page.get_pixmap(matrix=mat, alpha=False)
+            image = np.frombuffer(pixmap.samples, dtype=np.uint8).reshape(
+                pixmap.h, pixmap.w, pixmap.n
+            )
+            image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
+            if resize:
+                if img_size is None:
+                    img_size = (image.shape[1], image.shape[0])
+                else:
+                    if (image.shape[1], image.shape[0]) != img_size:
+                        image = cv2.resize(image, img_size)
+            images.append(image)
+    return images
+
+
+def file_to_images(
+    file_bytes: bytes,
+    file_type: Literal["IMAGE", "PDF"],
+    *,
+    max_img_size: Tuple[int, int],
+    max_num_imgs: int,
+) -> List[np.ndarray]:
+    if file_type == "IMAGE":
+        images = [image_bytes_to_array(file_bytes)]
+    elif file_type == "PDF":
+        images = read_pdf(file_bytes, resize=True, max_num_imgs=max_num_imgs)
+    else:
+        assert_never(file_type)
+    h, w = images[0].shape[0:2]
+    if w > max_img_size[1] or h > max_img_size[0]:
+        if w / h > max_img_size[0] / max_img_size[1]:
+            factor = max_img_size[0] / w
+        else:
+            factor = max_img_size[1] / h
+        images = [cv2.resize(img, (int(factor * w), int(factor * h))) for img in images]
+    return images
+
+
+def call_async(
+    func: Callable[_P, _R], /, *args: _P.args, **kwargs: _P.kwargs
+) -> Awaitable[_R]:
+    return asyncio.get_running_loop().run_in_executor(
+        None, partial(func, *args, **kwargs)
+    )

paddlex/inference/pipelines/single_model_pipeline.py

@@ -0,0 +1,76 @@
+# 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 .base import BasePipeline
+
+
+class _SingleModelPipeline(BasePipeline):
+
+    def __init__(self, model, batch_size=1, device=None, predictor_kwargs=None):
+        super().__init__(device, predictor_kwargs)
+        self._build_predictor(model)
+        self.set_predictor(batch_size=batch_size)
+
+    def _build_predictor(self, model):
+        self.model = self._create(model)
+
+    def set_predictor(self, batch_size=None, device=None):
+        if batch_size:
+            self.model.set_predictor(batch_size=batch_size)
+        if device:
+            self.model.set_predictor(device=device)
+
+    def predict(self, input, **kwargs):
+        self.set_predictor(**kwargs)
+        yield from self.model(input)
+
+
+class ImageClassification(_SingleModelPipeline):
+    entities = "image_classification"
+
+
+class ObjectDetection(_SingleModelPipeline):
+    entities = "object_detection"
+
+
+class InstanceSegmentation(_SingleModelPipeline):
+    entities = "instance_segmentation"
+
+
+class SemanticSegmentation(_SingleModelPipeline):
+    entities = "semantic_segmentation"
+
+
+class TSFc(_SingleModelPipeline):
+    entities = "ts_fc"
+
+
+class TSAd(_SingleModelPipeline):
+    entities = "ts_ad"
+
+
+class TSCls(_SingleModelPipeline):
+    entities = "ts_cls"
+
+
+class MultiLableImageClas(_SingleModelPipeline):
+    entities = "multi_label_image_classification"
+
+
+class SmallObjDet(_SingleModelPipeline):
+    entities = "small_object_detection"
+
+
+class AnomalyDetection(_SingleModelPipeline):
+    entities = "anomaly_detection"

paddlex/inference/pipelines/table_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 .table_recognition import _TableRecPipeline, TableRecPipeline

paddlex/inference/pipelines/table_recognition/table_recognition.py

@@ -0,0 +1,193 @@
+# 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 OCRResult, TableResult, StructureTableResult
+
+
+class _TableRecPipeline(BasePipeline):
+    """Table Recognition Pipeline"""
+
+    def __init__(
+        self,
+        device,
+        predictor_kwargs,
+    ):
+        super().__init__(device, predictor_kwargs)
+
+    def _build_predictor(
+        self,
+        layout_model,
+        text_det_model,
+        text_rec_model,
+        table_model,
+    ):
+        self.layout_predictor = self._create(model=layout_model)
+        self.ocr_pipeline = self._create(
+            pipeline=OCRPipeline,
+            text_det_model=text_det_model,
+            text_rec_model=text_rec_model,
+        )
+        self.table_predictor = self._create(model=table_model)
+        self._crop_by_boxes = CropByBoxes()
+        self._match = TableMatch(filter_ocr_result=False)
+
+    def set_predictor(
+        self,
+        layout_batch_size=None,
+        text_det_batch_size=None,
+        text_rec_batch_size=None,
+        table_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 device:
+            self.layout_predictor.set_predictor(device=device)
+            self.ocr_pipeline.text_rec_model.set_predictor(device=device)
+            self.table_predictor.set_predictor(device=device)
+
+    def predict(self, inputs):
+        raise NotImplementedError("The method `predict` has not been implemented yet.")
+
+    def get_related_ocr_result(self, box, ocr_res):
+        dt_polys_list = []
+        rec_text_list = []
+        score_list = []
+        unmatched_ocr_res = {"dt_polys": [], "rec_text": [], "rec_score": []}
+        unmatched_ocr_res["input_path"] = ocr_res["input_path"]
+        for i, text_box in enumerate(ocr_res["dt_polys"]):
+            text_box_area = convert_4point2rect(text_box)
+            if is_inside(text_box_area, box):
+                dt_polys_list.append(text_box)
+                rec_text_list.append(ocr_res["rec_text"][i])
+                score_list.append(ocr_res["rec_score"][i])
+            else:
+                unmatched_ocr_res["dt_polys"].append(text_box)
+                unmatched_ocr_res["rec_text"].append(ocr_res["rec_text"][i])
+                unmatched_ocr_res["rec_score"].append(ocr_res["rec_score"][i])
+        return (dt_polys_list, rec_text_list, score_list), unmatched_ocr_res
+
+    def get_table_result(self, input_imgs):
+        table_res_list = []
+        ocr_res_list = []
+        table_index = 0
+        img_list = [img["img"] for img in input_imgs]
+        for input_img, table_pred, ocr_pred in zip(
+            input_imgs, self.table_predictor(img_list), self.ocr_pipeline(img_list)
+        ):
+            single_table_box = table_pred["bbox"]
+            ori_x, ori_y, _, _ = input_img["box"]
+            ori_bbox_list = np.array(
+                get_ori_coordinate_for_table(ori_x, ori_y, single_table_box),
+                dtype=np.float32,
+            )
+            ori_ocr_bbox_list = np.array(
+                get_ori_coordinate_for_table(ori_x, ori_y, ocr_pred["dt_polys"]),
+                dtype=np.float32,
+            )
+            html_res = self._match(table_pred, ocr_pred)
+            ocr_pred["dt_polys"] = ori_ocr_bbox_list
+            table_res_list.append(
+                StructureTableResult(
+                    {
+                        "input_path": input_img["input_path"],
+                        "layout_bbox": [int(x) for x in input_img["box"]],
+                        "bbox": ori_bbox_list,
+                        "img_idx": table_index,
+                        "html": html_res,
+                    }
+                )
+            )
+            ocr_res_list.append(ocr_pred)
+            table_index += 1
+        return table_res_list, ocr_res_list
+
+
+class TableRecPipeline(_TableRecPipeline):
+    """Table Recognition Pipeline"""
+
+    entities = "table_recognition"
+
+    def __init__(
+        self,
+        layout_model,
+        text_det_model,
+        text_rec_model,
+        table_model,
+        layout_batch_size=1,
+        text_det_batch_size=1,
+        text_rec_batch_size=1,
+        table_batch_size=1,
+        device=None,
+        predictor_kwargs=None,
+    ):
+        super().__init__(device, predictor_kwargs)
+        self._build_predictor(layout_model, text_det_model, text_rec_model, table_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,
+        )
+
+    def predict(self, input, **kwargs):
+        self.set_predictor(**kwargs)
+        for layout_pred, ocr_pred in zip(
+            self.layout_predictor(input), self.ocr_pipeline(input)
+        ):
+            single_img_res = {
+                "input_path": "",
+                "layout_result": {},
+                "ocr_result": {},
+                "table_result": [],
+            }
+            # update layout result
+            single_img_res["input_path"] = layout_pred["input_path"]
+            single_img_res["layout_result"] = layout_pred
+            ocr_res = ocr_pred
+            table_subs = []
+            if len(layout_pred["boxes"]) > 0:
+                subs_of_img = list(self._crop_by_boxes(layout_pred))
+                # get cropped images with label "table"
+                for sub in subs_of_img:
+                    box = sub["box"]
+                    if sub["label"].lower() == "table":
+                        table_subs.append(sub)
+                        _, ocr_res = self.get_related_ocr_result(box, ocr_res)
+            table_res, all_table_ocr_res = self.get_table_result(table_subs)
+            for table_ocr_res in all_table_ocr_res:
+                ocr_res["dt_polys"].extend(table_ocr_res["dt_polys"])
+                ocr_res["rec_text"].extend(table_ocr_res["rec_text"])
+                ocr_res["rec_score"].extend(table_ocr_res["rec_score"])
+
+            single_img_res["table_result"] = table_res
+            single_img_res["ocr_result"] = OCRResult(ocr_res)
+
+            yield TableResult(single_img_res)