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

paddlex/inference/pipelines/serving/_pipeline_apps/face_recognition.py

@@ -0,0 +1,317 @@
+# 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 faiss
+import pickle
+from typing import Dict, List, Optional
+
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel, Field
+from typing_extensions import Annotated, TypeAlias
+
+from .....utils import logging
+from ....components.retrieval.faiss import IndexData
+from ...face_recognition import FaceRecPipeline
+from ..storage import create_storage
+from .. import utils as serving_utils
+from ..app import AppConfig, create_app
+from ..models import Response, ResultResponse
+
+
+class ImageLabelPair(BaseModel):
+    image: str
+    label: str
+
+
+class BuildIndexRequest(BaseModel):
+    imageLabelPairs: List[ImageLabelPair]
+
+
+class BuildIndexResult(BaseModel):
+    indexKey: str
+    idMap: Dict[int, str]
+
+
+class AddImagesToIndexRequest(BaseModel):
+    imageLabelPairs: List[ImageLabelPair]
+    indexKey: str
+
+
+class AddImagesToIndexResult(BaseModel):
+    idMap: Dict[int, str]
+
+
+class RemoveImagesFromIndexRequest(BaseModel):
+    ids: List[int]
+    indexKey: str
+
+
+class RemoveImagesFromIndexResult(BaseModel):
+    idMap: Dict[int, str]
+
+
+class InferRequest(BaseModel):
+    image: str
+    indexKey: Optional[str] = None
+
+
+BoundingBox: TypeAlias = Annotated[List[float], Field(min_length=4, max_length=4)]
+
+
+class RecResult(BaseModel):
+    label: str
+    score: float
+
+
+class Face(BaseModel):
+    bbox: BoundingBox
+    recResults: List[RecResult]
+    score: float
+
+
+class InferResult(BaseModel):
+    faces: List[Face]
+    image: str
+
+
+def _serialize_index_data(index_data: IndexData) -> bytes:
+    tup = (index_data.index_bytes, index_data.index_info)
+    return pickle.dumps(tup)
+
+
+def _deserialize_index_data(index_data_bytes: bytes) -> IndexData:
+    tup = pickle.loads(index_data_bytes)
+    index = faiss.deserialize_index(tup[0])
+    return IndexData(index, tup[1])
+
+
+def create_pipeline_app(pipeline: FaceRecPipeline, app_config: AppConfig) -> FastAPI:
+    app, ctx = create_app(
+        pipeline=pipeline, app_config=app_config, app_aiohttp_session=True
+    )
+
+    if ctx.config.extra and "index_storage" in ctx.config.extra:
+        ctx.extra["index_storage"] = create_storage(ctx.config.extra["index_storage"])
+    else:
+        ctx.extra["index_storage"] = create_storage({"type": "memory"})
+
+    @app.post(
+        "/face-recognition-index-build",
+        operation_id="buildIndex",
+        responses={422: {"model": Response}},
+    )
+    async def _build_index(
+        request: BuildIndexRequest,
+    ) -> ResultResponse[BuildIndexResult]:
+        pipeline = ctx.pipeline
+        aiohttp_session = ctx.aiohttp_session
+
+        request_id = serving_utils.generate_request_id()
+
+        try:
+            images = [pair.image for pair in request.imageLabelPairs]
+            file_bytes_list = await asyncio.gather(
+                *(serving_utils.get_raw_bytes(img, aiohttp_session) for img in images)
+            )
+            images = [
+                serving_utils.image_bytes_to_array(item) for item in file_bytes_list
+            ]
+            labels = [pair.label for pair in request.imageLabelPairs]
+
+            # TODO: Support specifying `index_type` and `metric_type` in the
+            # request
+            index_data = await pipeline.call(
+                pipeline.pipeline.build_index,
+                images,
+                labels,
+                index_type="Flat",
+                metric_type="IP",
+            )
+
+            index_storage = ctx.extra["index_storage"]
+            index_key = request_id
+            index_data_bytes = await serving_utils.call_async(
+                _serialize_index_data, index_data
+            )
+            await serving_utils.call_async(
+                index_storage.set, index_key, index_data_bytes
+            )
+
+            return ResultResponse(
+                logId=serving_utils.generate_log_id(),
+                errorCode=0,
+                errorMsg="Success",
+                result=BuildIndexResult(indexKey=index_key, idMap=index_data.id_map),
+            )
+
+        except Exception as e:
+            logging.exception(e)
+            raise HTTPException(status_code=500, detail="Internal server error")
+
+    @app.post(
+        "/face-recognition-index-add",
+        operation_id="buildIndex",
+        responses={422: {"model": Response}},
+    )
+    async def _add_images_to_index(
+        request: AddImagesToIndexRequest,
+    ) -> ResultResponse[AddImagesToIndexResult]:
+        pipeline = ctx.pipeline
+        aiohttp_session = ctx.aiohttp_session
+
+        try:
+            images = [pair.image for pair in request.imageLabelPairs]
+            file_bytes_list = await asyncio.gather(
+                *(serving_utils.get_raw_bytes(img, aiohttp_session) for img in images)
+            )
+            images = [
+                serving_utils.image_bytes_to_array(item) for item in file_bytes_list
+            ]
+            labels = [pair.label for pair in request.imageLabelPairs]
+            index_storage = ctx.extra["index_storage"]
+            index_data_bytes = await serving_utils.call_async(
+                index_storage.get, request.indexKey
+            )
+            index_data = await serving_utils.call_async(
+                _deserialize_index_data, index_data_bytes
+            )
+
+            index_data = await pipeline.call(
+                pipeline.pipeline.append_index, images, labels, index_data
+            )
+
+            index_data_bytes = await serving_utils.call_async(
+                _serialize_index_data, index_data
+            )
+            await serving_utils.call_async(
+                index_storage.set, request.indexKey, index_data_bytes
+            )
+
+            return ResultResponse(
+                logId=serving_utils.generate_log_id(),
+                errorCode=0,
+                errorMsg="Success",
+                result=AddImagesToIndexResult(idMap=index_data.id_map),
+            )
+
+        except Exception as e:
+            logging.exception(e)
+            raise HTTPException(status_code=500, detail="Internal server error")
+
+    @app.post(
+        "/face-recognition-index-remove",
+        operation_id="buildIndex",
+        responses={422: {"model": Response}},
+    )
+    async def _remove_images_from_index(
+        request: RemoveImagesFromIndexRequest,
+    ) -> ResultResponse[RemoveImagesFromIndexResult]:
+        pipeline = ctx.pipeline
+
+        try:
+            index_storage = ctx.extra["index_storage"]
+            index_data_bytes = await serving_utils.call_async(
+                index_storage.get, request.indexKey
+            )
+            index_data = await serving_utils.call_async(
+                _deserialize_index_data, index_data_bytes
+            )
+
+            index_data = await pipeline.call(
+                pipeline.pipeline.remove_index, request.ids, index_data
+            )
+
+            index_data_bytes = await serving_utils.call_async(
+                _serialize_index_data, index_data
+            )
+            await serving_utils.call_async(
+                index_storage.set, request.indexKey, index_data_bytes
+            )
+
+            return ResultResponse(
+                logId=serving_utils.generate_log_id(),
+                errorCode=0,
+                errorMsg="Success",
+                result=RemoveImagesFromIndexResult(idMap=index_data.id_map),
+            )
+
+        except Exception as e:
+            logging.exception(e)
+            raise HTTPException(status_code=500, detail="Internal server error")
+
+    @app.post(
+        "/face-recognition-infer",
+        operation_id="infer",
+        responses={422: {"model": Response}},
+    )
+    async def _infer(request: InferRequest) -> ResultResponse[InferResult]:
+        pipeline = ctx.pipeline
+        aiohttp_session = ctx.aiohttp_session
+
+        try:
+            image_bytes = await serving_utils.get_raw_bytes(
+                request.image, aiohttp_session
+            )
+            image = serving_utils.image_bytes_to_array(image_bytes)
+
+            if request.indexKey is not None:
+                index_storage = ctx.extra["index_storage"]
+                index_data_bytes = await serving_utils.call_async(
+                    index_storage.get, request.indexKey
+                )
+                index_data = await serving_utils.call_async(
+                    _deserialize_index_data, index_data_bytes
+                )
+            else:
+                index_data = None
+
+            result = list(
+                await pipeline.call(pipeline.pipeline.predict, image, index=index_data)
+            )[0]
+
+            faces: List[Face] = []
+            for face in result["boxes"]:
+                rec_results: List[RecResult] = []
+                if face["rec_scores"] is not None:
+                    for label, score in zip(face["labels"], face["rec_scores"]):
+                        rec_results.append(
+                            RecResult(
+                                label=label,
+                                score=score,
+                            )
+                        )
+                faces.append(
+                    Face(
+                        bbox=face["coordinate"],
+                        recResults=rec_results,
+                        score=face["det_score"],
+                    )
+                )
+            output_image_base64 = serving_utils.base64_encode(
+                serving_utils.image_to_bytes(result.img)
+            )
+
+            return ResultResponse(
+                logId=serving_utils.generate_log_id(),
+                errorCode=0,
+                errorMsg="Success",
+                result=InferResult(faces=faces, image=output_image_base64),
+            )
+
+        except Exception as e:
+            logging.exception(e)
+            raise HTTPException(status_code=500, detail="Internal server error")
+
+    return app

paddlex/inference/pipelines/serving/_pipeline_apps/formula_recognition.py

@@ -0,0 +1,119 @@
+# 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 List, Optional
+
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel, Field
+from typing_extensions import Annotated, TypeAlias
+
+from .....utils import logging
+from ...formula_recognition import FormulaRecognitionPipeline
+from .. import utils as serving_utils
+from ..app import AppConfig, create_app
+from ..models import Response, ResultResponse
+
+
+class InferenceParams(BaseModel):
+    maxLongSide: Optional[Annotated[int, Field(gt=0)]] = None
+
+
+class InferRequest(BaseModel):
+    image: str
+    inferenceParams: Optional[InferenceParams] = None
+
+
+Point: TypeAlias = Annotated[List[float], Field(min_length=2, max_length=2)]
+Polygon: TypeAlias = Annotated[List[Point], Field(min_length=3)]
+
+
+class Formula(BaseModel):
+    poly: Polygon
+    latex: str
+
+
+class InferResult(BaseModel):
+    formulas: List[Formula]
+    layoutImage: str
+    ocrImage: Optional[str] = None
+
+
+def create_pipeline_app(
+    pipeline: FormulaRecognitionPipeline, app_config: AppConfig
+) -> FastAPI:
+    app, ctx = create_app(
+        pipeline=pipeline, app_config=app_config, app_aiohttp_session=True
+    )
+
+    @app.post(
+        "/formula-recognition",
+        operation_id="infer",
+        responses={422: {"model": Response}},
+        response_model_exclude_none=True,
+    )
+    async def _infer(request: InferRequest) -> ResultResponse[InferResult]:
+        pipeline = ctx.pipeline
+        aiohttp_session = ctx.aiohttp_session
+
+        if request.inferenceParams:
+            max_long_side = request.inferenceParams.maxLongSide
+            if max_long_side:
+                raise HTTPException(
+                    status_code=422,
+                    detail="`max_long_side` is currently not supported.",
+                )
+
+        try:
+            file_bytes = await serving_utils.get_raw_bytes(
+                request.image, aiohttp_session
+            )
+            image = serving_utils.image_bytes_to_array(file_bytes)
+
+            result = (await pipeline.infer(image))[0]
+
+            formulas: List[Formula] = []
+            for poly, latex in zip(result["dt_polys"], result["rec_formula"]):
+                formulas.append(
+                    Formula(
+                        poly=poly,
+                        latex=latex,
+                    )
+                )
+            layout_image_base64 = serving_utils.base64_encode(
+                serving_utils.image_to_bytes(result["layout_result"].img)
+            )
+            ocr_image = result["formula_result"].img
+            if ocr_image is not None:
+                ocr_image_base64 = serving_utils.base64_encode(
+                    serving_utils.image_to_bytes(ocr_image)
+                )
+            else:
+                ocr_image_base64 = None
+
+            return ResultResponse(
+                logId=serving_utils.generate_log_id(),
+                errorCode=0,
+                errorMsg="Success",
+                result=InferResult(
+                    formulas=formulas,
+                    layoutImage=layout_image_base64,
+                    ocrImage=ocr_image_base64,
+                ),
+            )
+
+        except Exception as e:
+            logging.exception(e)
+            raise HTTPException(status_code=500, detail="Internal server error")
+
+    return app

paddlex/inference/pipelines/serving/_pipeline_apps/image_classification.py

@@ -0,0 +1,101 @@
+# 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 itertools import islice
+from typing import List, Optional
+
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel, Field
+from typing_extensions import Annotated
+
+from .....utils import logging
+from ...single_model_pipeline import ImageClassification
+from .. import utils as serving_utils
+from ..app import AppConfig, create_app
+from ..models import Response, ResultResponse
+
+
+class InferenceParams(BaseModel):
+    topK: Optional[Annotated[int, Field(gt=0)]] = None
+
+
+class InferRequest(BaseModel):
+    image: str
+    inferenceParams: Optional[InferenceParams] = None
+
+
+class Category(BaseModel):
+    id: int
+    name: str
+    score: float
+
+
+class InferResult(BaseModel):
+    categories: List[Category]
+    image: str
+
+
+def create_pipeline_app(
+    pipeline: ImageClassification, app_config: AppConfig
+) -> FastAPI:
+    app, ctx = create_app(
+        pipeline=pipeline, app_config=app_config, app_aiohttp_session=True
+    )
+
+    @app.post(
+        "/image-classification",
+        operation_id="infer",
+        responses={422: {"model": Response}},
+    )
+    async def _infer(request: InferRequest) -> ResultResponse[InferResult]:
+        pipeline = ctx.pipeline
+        aiohttp_session = ctx.aiohttp_session
+
+        try:
+            file_bytes = await serving_utils.get_raw_bytes(
+                request.image, aiohttp_session
+            )
+            image = serving_utils.image_bytes_to_array(file_bytes)
+            top_k: Optional[int] = None
+            if request.inferenceParams is not None:
+                if request.inferenceParams.topK is not None:
+                    top_k = request.inferenceParams.topK
+
+            result = (await pipeline.infer(image))[0]
+
+            if "label_names" in result:
+                cat_names = result["label_names"]
+            else:
+                cat_names = [str(id_) for id_ in result["class_ids"]]
+            categories: List[Category] = []
+            for id_, name, score in islice(
+                zip(result["class_ids"], cat_names, result["scores"]), None, top_k
+            ):
+                categories.append(Category(id=id_, name=name, score=score))
+            output_image_base64 = serving_utils.base64_encode(
+                serving_utils.image_to_bytes(result.img)
+            )
+
+            return ResultResponse(
+                logId=serving_utils.generate_log_id(),
+                errorCode=0,
+                errorMsg="Success",
+                result=InferResult(categories=categories, image=output_image_base64),
+            )
+
+        except Exception as e:
+            logging.exception(e)
+            raise HTTPException(status_code=500, detail="Internal server error")
+
+    return app

paddlex/inference/pipelines/serving/_pipeline_apps/instance_segmentation.py

@@ -0,0 +1,112 @@
+# 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 List
+
+import numpy as np
+import pycocotools.mask as mask_util
+from fastapi import FastAPI, HTTPException
+from pydantic import BaseModel, Field
+from typing_extensions import Annotated, TypeAlias
+
+from .....utils import logging
+from ...single_model_pipeline import InstanceSegmentation
+from .. import utils as serving_utils
+from ..app import AppConfig, create_app
+from ..models import Response, ResultResponse
+
+
+class InferRequest(BaseModel):
+    image: str
+
+
+BoundingBox: TypeAlias = Annotated[List[float], Field(min_length=4, max_length=4)]
+
+
+class Mask(BaseModel):
+    rleResult: str
+    size: Annotated[List[int], Field(min_length=2, max_length=2)]
+
+
+class Instance(BaseModel):
+    bbox: BoundingBox
+    categoryId: int
+    score: float
+    mask: Mask
+
+
+class InferResult(BaseModel):
+    instances: List[Instance]
+    image: str
+
+
+def _rle(mask: np.ndarray) -> str:
+    rle_res = mask_util.encode(np.asarray(mask[..., None], order="F", dtype="uint8"))[0]
+    return rle_res["counts"].decode("utf-8")
+
+
+def create_pipeline_app(
+    pipeline: InstanceSegmentation, app_config: AppConfig
+) -> FastAPI:
+    app, ctx = create_app(
+        pipeline=pipeline,
+        app_config=app_config,
+        app_aiohttp_session=True,
+    )
+
+    @app.post(
+        "/instance-segmentation",
+        operation_id="infer",
+        responses={422: {"model": Response}},
+    )
+    async def _infer(request: InferRequest) -> ResultResponse[InferResult]:
+        pipeline = ctx.pipeline
+        aiohttp_session = ctx.aiohttp_session
+
+        try:
+            file_bytes = await serving_utils.get_raw_bytes(
+                request.image, aiohttp_session
+            )
+            image = serving_utils.image_bytes_to_array(file_bytes)
+
+            result = (await pipeline.infer(image))[0]
+
+            instances: List[Instance] = []
+            for obj, mask in zip(result["boxes"], result["masks"]):
+                rle_res = _rle(mask)
+                mask = Mask(rleResult=rle_res, size=mask.shape)
+                instances.append(
+                    Instance(
+                        bbox=obj["coordinate"],
+                        categoryId=obj["cls_id"],
+                        score=obj["score"],
+                        mask=mask,
+                    )
+                )
+            output_image_base64 = serving_utils.base64_encode(
+                serving_utils.image_to_bytes(result.img)
+            )
+
+            return ResultResponse(
+                logId=serving_utils.generate_log_id(),
+                errorCode=0,
+                errorMsg="Success",
+                result=InferResult(instances=instances, image=output_image_base64),
+            )
+
+        except Exception as e:
+            logging.exception(e)
+            raise HTTPException(status_code=500, detail="Internal server error")
+
+    return app