hpi refactor (#2688)

libs/paddlex-hpi/requirements.txt

@@ -1,4 +1,4 @@
-# ultrainfer
+# ultra_infer
 # paddlex
 importlib-resources >= 6.4
 numpy >= 1.21

libs/paddlex-hpi/src/paddlex_hpi/_config.py

@@ -16,7 +16,7 @@ import warnings
 from pathlib import Path
 from typing import Any, Dict, List, Mapping, Optional, Tuple, Type, Union
 
-import ultrainfer as ui
+import ultra_infer as ui
 from paddlex.utils import logging
 from pydantic import BaseModel, ConfigDict, Field, field_validator
 from typing_extensions import Annotated, TypeAlias, TypedDict, assert_never

libs/paddlex-hpi/src/paddlex_hpi/models/anomaly_detection.py

@@ -12,14 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any, List
+from typing import Any, Dict, List
 
-import ultrainfer as ui
+import ultra_infer as ui
 import numpy as np
+from paddlex.inference.common.batch_sampler import ImageBatchSampler
 from paddlex.inference.results import SegResult
 from paddlex.modules.anomaly_detection.model_list import MODELS
 
-from paddlex_hpi._utils.typing import BatchData, Data
 from paddlex_hpi.models.base import CVPredictor
 
 
@@ -37,20 +37,27 @@ class UadPredictor(CVPredictor):
         )
         return model
 
-    def _predict(self, batch_data: BatchData) -> BatchData:
-        imgs = [np.ascontiguousarray(data["img"]) for data in batch_data]
+    def _build_batch_sampler(self) -> ImageBatchSampler:
+        return ImageBatchSampler()
+
+    def _get_result_class(self) -> type:
+        return SegResult
+
+    def process(self, batch_data: List[Any]) -> Dict[str, List[Any]]:
+        batch_raw_imgs = self._data_reader(imgs=batch_data)
+        imgs = [np.ascontiguousarray(img) for img in batch_raw_imgs]
         ui_results = self._ui_model.batch_predict(imgs)
-        results: BatchData = []
-        for data, ui_result in zip(batch_data, ui_results):
-            uad_result = self._create_uad_result(data, ui_result)
-            results.append({"result": uad_result})
-        return results
-
-    def _create_uad_result(self, data: Data, ui_result: Any) -> SegResult:
-        pred = np.array(ui_result.label_map, dtype=np.int32).reshape(ui_result.shape)
-        pred = pred[np.newaxis]
-        dic = {
-            "input_path": data["input_path"],
-            "pred": pred,
+
+        preds_list = []
+        for ui_result in ui_results:
+            pred = np.array(ui_result.label_map, dtype=np.int32).reshape(
+                ui_result.shape
+            )
+            pred = pred[np.newaxis]
+            preds_list.append(pred)
+
+        return {
+            "input_path": batch_data,
+            "input_img": batch_raw_imgs,
+            "pred": preds_list,
         }
-        return SegResult(dic)

libs/paddlex-hpi/src/paddlex_hpi/models/base.py

@@ -19,18 +19,16 @@ from typing import (
     Any,
     Dict,
     Final,
-    Generator,
-    List,
+    Iterator,
     Optional,
-    Protocol,
     TypedDict,
     Union,
 )
 
-import ultrainfer as ui
-from ultrainfer.model import BaseUltraInferModel
-from paddlex.inference.components import ReadImage, ReadTS
-from paddlex.inference.models import BasePredictor
+import ultra_infer as ui
+from ultra_infer.model import BaseUltraInferModel
+from paddlex.inference.common.reader import ReadImage
+from paddlex.inference.models_new import BasePredictor
 from paddlex.inference.utils.new_ir_blacklist import NEWIR_BLOCKLIST
 from paddlex.utils import device as device_helper
 from paddlex.utils import logging
@@ -38,7 +36,7 @@ from paddlex.utils.subclass_register import AutoRegisterABCMetaClass
 from typing_extensions import assert_never
 
 from paddlex_hpi._config import HPIConfig
-from paddlex_hpi._utils.typing import Backend, BatchData
+from paddlex_hpi._utils.typing import Backend
 
 HPI_CONFIG_KEY: Final[str] = "Hpi"
 
@@ -64,6 +62,11 @@ class HPPredictor(BasePredictor, metaclass=AutoRegisterABCMetaClass):
         self._hpi_params = hpi_params or {}
         self._hpi_config = self._get_hpi_config()
         self._ui_model = self.build_ui_model()
+        self._data_reader = self._build_data_reader()
+
+    def __call__(self, input: Any, **kwargs: dict[str, Any]) -> Iterator[Any]:
+        self.set_predictor(**kwargs)
+        yield from self.apply(input)
 
     @property
     def model_path(self) -> Path:
@@ -79,6 +82,8 @@ class HPPredictor(BasePredictor, metaclass=AutoRegisterABCMetaClass):
             if device is not None:
                 if device != self._device:
                     raise RuntimeError("Currently, changing devices is not supported.")
+        if "batch_size" in kwargs:
+            self.batch_sampler.batch_size = kwargs.pop("batch_size")
         if kwargs:
             raise TypeError(f"Unexpected arguments: {kwargs}")
 
@@ -86,10 +91,6 @@ class HPPredictor(BasePredictor, metaclass=AutoRegisterABCMetaClass):
         option = self._create_ui_option()
         return self._build_ui_model(option)
 
-    @abc.abstractmethod
-    def _build_ui_model(self, option: ui.RuntimeOption) -> BaseUltraInferModel:
-        raise NotImplementedError
-
     def _get_hpi_config(self) -> HPIConfig:
         if HPI_CONFIG_KEY not in self.config:
             logging.debug("Key %r not found in the config", HPI_CONFIG_KEY)
@@ -134,56 +135,20 @@ class HPPredictor(BasePredictor, metaclass=AutoRegisterABCMetaClass):
         backend_config.update_ui_option(option, self.model_dir)
         return option
 
-
-class _DataReaderLike(Protocol):
-    batch_size: int
-
-    def __call__(self, input_list: Any) -> Generator[BatchData, None, None]: ...
-
-
-class HPPredictorWithDataReader(HPPredictor):
-    def __init__(
-        self,
-        model_dir: Union[str, PathLike],
-        config: Optional[Dict[str, Any]] = None,
-        device: Optional[str] = None,
-        hpi_params: Optional[HPIParams] = None,
-    ) -> None:
-        super().__init__(
-            model_dir=model_dir,
-            config=config,
-            device=device,
-            hpi_params=hpi_params,
-        )
-        self._batch_size = 1
-        self._data_reader = self._build_data_reader()
-
-    def set_predictor(self, **kwargs: Any) -> None:
-        batch_size = kwargs.pop("batch_size", None)
-        super().set_predictor(**kwargs)
-        if batch_size is not None:
-            self._batch_size = batch_size
-            self._data_reader.batch_size = batch_size
-            logging.info("Batch size updated to %d", self._batch_size)
-
-    def apply(self, input: Any) -> Generator[BatchData, None, None]:
-        for batch_data in self._data_reader(input):
-            yield self._predict(batch_data)
-
     @abc.abstractmethod
-    def _build_data_reader(self) -> _DataReaderLike:
+    def _build_ui_model(self, option: ui.RuntimeOption) -> BaseUltraInferModel:
         raise NotImplementedError
 
     @abc.abstractmethod
-    def _predict(self, batch_data: BatchData) -> BatchData:
+    def _build_data_reader(self):
         raise NotImplementedError
 
 
-class CVPredictor(HPPredictorWithDataReader):
-    def _build_data_reader(self) -> _DataReaderLike:
-        return ReadImage(batch_size=self._batch_size, format="BGR")
+class CVPredictor(HPPredictor):
+    def _build_data_reader(self):
+        return ReadImage(format="BGR")
 
 
-class TSPredictor(HPPredictorWithDataReader):
-    def _build_data_reader(self) -> _DataReaderLike:
-        return ReadTS(batch_size=self._batch_size)
+class TSPredictor(HPPredictor):
+    def _build_data_reader(self):
+        return None

libs/paddlex-hpi/src/paddlex_hpi/models/formula_recognition.py

@@ -12,14 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any, List
+from typing import Any, Dict, List
 
-import ultrainfer as ui
+import ultra_infer as ui
 import numpy as np
+from paddlex.inference.common.batch_sampler import ImageBatchSampler
 from paddlex.inference.results import FormulaRecResult
 from paddlex.modules.formula_recognition.model_list import MODELS
 
-from paddlex_hpi._utils.typing import BatchData, Data
 from paddlex_hpi.models.base import CVPredictor
 
 
@@ -37,20 +37,23 @@ class LaTeXOCRPredictor(CVPredictor):
         )
         return model
 
-    def _predict(self, batch_data: BatchData) -> BatchData:
-        imgs = [
-            np.ascontiguousarray(data["img"]).astype("float32") for data in batch_data
-        ]
+    def _build_batch_sampler(self) -> ImageBatchSampler:
+        return ImageBatchSampler()
+
+    def _get_result_class(self) -> type:
+        return FormulaRecResult
+
+    def process(self, batch_data: List[Any]) -> Dict[str, List[Any]]:
+        batch_raw_imgs = self._data_reader(imgs=batch_data)
+        imgs = [np.ascontiguousarray(img) for img in batch_raw_imgs]
         ui_results = self._ui_model.batch_predict(imgs)
-        results: BatchData = []
-        for data, ui_result in zip(batch_data, ui_results):
-            rec_result = self._create_rec_result(data, ui_result)
-            results.append({"result": rec_result})
-        return results
-
-    def _create_rec_result(self, data: Data, ui_result: Any) -> FormulaRecResult:
-        dic = {
-            "input_path": data["input_path"],
-            "rec_text": ui_result.rec_text,
+
+        rec_text_list = []
+        for ui_result in ui_results:
+            rec_text_list.append(ui_result.rec_text)
+
+        return {
+            "input_path": batch_data,
+            "input_img": batch_raw_imgs,
+            "rec_text": rec_text_list,
         }
-        return FormulaRecResult(dic)

libs/paddlex-hpi/src/paddlex_hpi/models/general_recognition.py

@@ -12,14 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any, List
+from typing import Any, Dict, List
 
-import ultrainfer as ui
+import ultra_infer as ui
 import numpy as np
+from paddlex.inference.common.batch_sampler import ImageBatchSampler
 from paddlex.inference.results import BaseResult
 from paddlex.modules.general_recognition.model_list import MODELS
 
-from paddlex_hpi._utils.typing import BatchData, Data
 from paddlex_hpi.models.base import CVPredictor
 
 
@@ -37,20 +37,23 @@ class ShiTuRecPredictor(CVPredictor):
         )
         return model
 
-    def _predict(self, batch_data: BatchData) -> BatchData:
-        imgs = [
-            np.ascontiguousarray(data["img"]).astype("float32") for data in batch_data
-        ]
+    def _build_batch_sampler(self) -> ImageBatchSampler:
+        return ImageBatchSampler()
+
+    def _get_result_class(self) -> type:
+        return BaseResult
+
+    def process(self, batch_data: List[Any]) -> Dict[str, List[Any]]:
+        batch_raw_imgs = self._data_reader(imgs=batch_data)
+        imgs = [np.ascontiguousarray(img) for img in batch_raw_imgs]
         ui_results = self._ui_model.batch_predict(imgs)
-        results: BatchData = []
-        for data, ui_result in zip(batch_data, ui_results):
-            clas_result = self._create_rec_result(data, ui_result)
-            results.append({"result": clas_result})
-        return results
-
-    def _create_rec_result(self, data: Data, ui_result: Any) -> BaseResult:
-        dic = {
-            "input_path": data["input_path"],
-            "feature": np.array(ui_result.feature, dtype="float32"),
+
+        feature_list = []
+        for ui_result in ui_results:
+            feature_list.append(np.array(ui_result.feature, dtype="float32"))
+
+        return {
+            "input_path": batch_data,
+            "input_img": batch_raw_imgs,
+            "feature": feature_list,
         }
-        return BaseResult(dic)

libs/paddlex-hpi/src/paddlex_hpi/models/image_classification.py

@@ -15,13 +15,13 @@
 import os
 from typing import Any, Dict, List, Optional, Union
 
-import ultrainfer as ui
+import ultra_infer as ui
 import numpy as np
+from paddlex.inference.common.batch_sampler import ImageBatchSampler
 from paddlex.inference.results import TopkResult
 from paddlex.modules.image_classification.model_list import MODELS
 from pydantic import BaseModel
 
-from paddlex_hpi._utils.typing import BatchData, Data
 from paddlex_hpi.models.base import CVPredictor, HPIParams
 
 
@@ -49,6 +49,12 @@ class ClasPredictor(CVPredictor):
         self._pp_params = self._get_pp_params()
         self._ui_model.postprocessor.topk = self._pp_params.topk
 
+    def _build_batch_sampler(self) -> ImageBatchSampler:
+        return ImageBatchSampler()
+
+    def _get_result_class(self) -> type:
+        return TopkResult
+
     def _build_ui_model(
         self, option: ui.RuntimeOption
     ) -> ui.vision.classification.PaddleClasModel:
@@ -60,14 +66,29 @@ class ClasPredictor(CVPredictor):
         )
         return model
 
-    def _predict(self, batch_data: BatchData) -> BatchData:
-        imgs = [np.ascontiguousarray(data["img"]) for data in batch_data]
+    def process(self, batch_data: List[Any]) -> Dict[str, List[Any]]:
+        batch_raw_imgs = self._data_reader(imgs=batch_data)
+        imgs = [np.ascontiguousarray(img) for img in batch_raw_imgs]
         ui_results = self._ui_model.batch_predict(imgs)
-        results: BatchData = []
-        for data, ui_result in zip(batch_data, ui_results):
-            clas_result = self._create_clas_result(data, ui_result)
-            results.append({"result": clas_result})
-        return results
+
+        class_ids_list = []
+        scores_list = []
+        label_names_list = []
+        for ui_result in ui_results:
+            class_ids_list.append(ui_result.label_ids)
+            scores_list.append(np.around(ui_result.scores, decimals=5).tolist())
+            if self._pp_params.label_list is not None:
+                label_names_list.append(
+                    [self._pp_params.label_list[i] for i in ui_result.label_ids]
+                )
+
+        return {
+            "input_path": batch_data,
+            "input_img": batch_raw_imgs,
+            "class_ids": class_ids_list,
+            "scores": scores_list,
+            "label_names": label_names_list,
+        }
 
     def _get_pp_params(self) -> _ClasPPParams:
         pp_config = self.config["PostProcess"]
@@ -77,15 +98,3 @@ class ClasPredictor(CVPredictor):
         topk = topk_config["topk"]
         label_list = topk_config.get("label_list", None)
         return _ClasPPParams(topk=topk, label_list=label_list)
-
-    def _create_clas_result(self, data: Data, ui_result: Any) -> TopkResult:
-        dic = {
-            "input_path": data["input_path"],
-            "class_ids": ui_result.label_ids,
-            "scores": np.around(ui_result.scores, decimals=5).tolist(),
-        }
-        if self._pp_params.label_list is not None:
-            dic["label_names"] = [
-                self._pp_params.label_list[i] for i in ui_result.label_ids
-            ]
-        return TopkResult(dic)

libs/paddlex-hpi/src/paddlex_hpi/models/image_unwarping.py

@@ -12,14 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any, List
+from typing import Any, Dict, List
 
-import ultrainfer as ui
+import ultra_infer as ui
 import numpy as np
+from paddlex.inference.common.batch_sampler import ImageBatchSampler
 from paddlex.inference.results import DocTrResult
 from paddlex.modules.image_unwarping.model_list import MODELS
 
-from paddlex_hpi._utils.typing import BatchData, Data
 from paddlex_hpi.models.base import CVPredictor
 
 
@@ -34,23 +34,28 @@ class WarpPredictor(CVPredictor):
         )
         return model
 
-    def _predict(self, batch_data: BatchData) -> BatchData:
-        imgs = [np.ascontiguousarray(data["img"]) for data in batch_data]
+    def _build_batch_sampler(self) -> ImageBatchSampler:
+        return ImageBatchSampler()
+
+    def _get_result_class(self) -> type:
+        return DocTrResult
+
+    def process(self, batch_data: List[Any]) -> Dict[str, List[Any]]:
+        batch_raw_imgs = self._data_reader(imgs=batch_data)
+        imgs = [np.ascontiguousarray(img) for img in batch_raw_imgs]
         ui_results = self._ui_model.batch_predict(imgs)
-        results: BatchData = []
-        for data, ui_result in zip(batch_data, ui_results):
-            warp_result = self._create_warp_result(data, ui_result)
-            results.append({"result": warp_result})
-        return results
-
-    def _create_warp_result(self, data: Data, ui_result: Any) -> DocTrResult:
-        img = ui_result.numpy()
-        img = np.moveaxis(img[0], 0, 2)
-        img *= 255
-        img = img[:, :, ::-1]
-        img = img.astype("uint8")
-        dic = {
-            "input_path": data["input_path"],
-            "doctr_img": img,
+
+        doctr_img_list = []
+        for ui_result in ui_results:
+            img = ui_result.numpy()
+            img = np.moveaxis(img[0], 0, 2)
+            img *= 255
+            img = img[:, :, ::-1]
+            img = img.astype("uint8")
+            doctr_img_list.append(img)
+
+        return {
+            "input_path": batch_data,
+            "input_img": batch_raw_imgs,
+            "doctr_img": doctr_img_list,
         }
-        return DocTrResult(dic)

libs/paddlex-hpi/src/paddlex_hpi/models/instance_segmentation.py

@@ -15,13 +15,13 @@
 import os
 from typing import Any, Dict, List, Optional, Union
 
-import ultrainfer as ui
+import ultra_infer as ui
 import numpy as np
+from paddlex.inference.common.batch_sampler import ImageBatchSampler
 from paddlex.inference.results import InstanceSegResult
 from paddlex.modules.instance_segmentation.model_list import MODELS
 from pydantic import BaseModel
 
-from paddlex_hpi._utils.typing import BatchData, Data
 from paddlex_hpi.models.base import CVPredictor, HPIParams
 
 
@@ -59,39 +59,36 @@ class InstanceSegPredictor(CVPredictor):
         )
         return model
 
-    def _predict(self, batch_data: BatchData) -> BatchData:
-        imgs = [np.ascontiguousarray(data["img"]) for data in batch_data]
-        ui_results = self._ui_model.batch_predict(imgs)
-        results: BatchData = []
-        for data, ui_result in zip(batch_data, ui_results):
-            instance_seg_result = self._create_instance_seg_result(data, ui_result)
-            results.append({"result": instance_seg_result})
-        return results
+    def _build_batch_sampler(self) -> ImageBatchSampler:
+        return ImageBatchSampler()
 
-    def _get_pp_params(self) -> _InstanceSegPPParams:
-        return _InstanceSegPPParams(
-            threshold=self.config["draw_threshold"],
-            label_list=self.config["label_list"],
-        )
+    def _get_result_class(self) -> type:
+        return InstanceSegResult
 
-    def _create_instance_seg_result(
-        self, data: Data, ui_result: Any
-    ) -> InstanceSegResult:
-        inds = sorted(
-            range(len(ui_result.scores)), key=ui_result.scores.__getitem__, reverse=True
-        )
-        inds = [i for i in inds if ui_result.scores[i] > self._pp_params.threshold]
-        inds = [i for i in inds if ui_result.label_ids[i] > -1]
-        ids = [ui_result.label_ids[i] for i in inds]
-        scores = [ui_result.scores[i] for i in inds]
-        boxes = [ui_result.boxes[i] for i in inds]
-        masks = [ui_result.masks[i] for i in inds]
-        masks = [
-            np.array(mask.data, dtype=np.uint8).reshape(mask.shape) for mask in masks
-        ]
-        dic = {
-            "input_path": data["input_path"],
-            "boxes": [
+    def process(self, batch_data: List[Any]) -> Dict[str, List[Any]]:
+        batch_raw_imgs = self._data_reader(imgs=batch_data)
+        imgs = [np.ascontiguousarray(img) for img in batch_raw_imgs]
+        ui_results = self._ui_model.batch_predict(imgs)
+
+        boxes_list = []
+        masks_list = []
+        for ui_result in ui_results:
+            inds = sorted(
+                range(len(ui_result.scores)),
+                key=ui_result.scores.__getitem__,
+                reverse=True,
+            )
+            inds = [i for i in inds if ui_result.scores[i] > self._pp_params.threshold]
+            inds = [i for i in inds if ui_result.label_ids[i] > -1]
+            ids = [ui_result.label_ids[i] for i in inds]
+            scores = [ui_result.scores[i] for i in inds]
+            boxes = [ui_result.boxes[i] for i in inds]
+            masks = [ui_result.masks[i] for i in inds]
+            masks = [
+                np.array(mask.data, dtype=np.uint8).reshape(mask.shape)
+                for mask in masks
+            ]
+            boxes = [
                 {
                     "cls_id": id_,
                     "label": self._pp_params.label_list[id_],
@@ -99,7 +96,19 @@ class InstanceSegPredictor(CVPredictor):
                     "coordinate": box,
                 }
                 for id_, score, box in zip(ids, scores, boxes)
-            ],
-            "masks": masks,
+            ]
+            boxes_list.append(boxes)
+            masks_list.append(masks)
+
+        return {
+            "input_path": batch_data,
+            "input_img": batch_raw_imgs,
+            "boxes": boxes_list,
+            "masks": masks_list,
         }
-        return InstanceSegResult(dic)
+
+    def _get_pp_params(self) -> _InstanceSegPPParams:
+        return _InstanceSegPPParams(
+            threshold=self.config["draw_threshold"],
+            label_list=self.config["label_list"],
+        )

libs/paddlex-hpi/src/paddlex_hpi/models/multilabel_classification.py

@@ -15,12 +15,12 @@
 import os
 from typing import Any, Dict, List, Optional, Union
 
-import ultrainfer as ui
+import ultra_infer as ui
 import numpy as np
+from paddlex.inference.common.batch_sampler import ImageBatchSampler
 from paddlex.inference.results import MLClassResult
 from paddlex.modules.multilabel_classification.model_list import MODELS
 
-from paddlex_hpi._utils.typing import BatchData, Data
 from paddlex_hpi.models.base import CVPredictor, HPIParams
 
 
@@ -53,14 +53,35 @@ class MLClasPredictor(CVPredictor):
         )
         return model
 
-    def _predict(self, batch_data: BatchData) -> BatchData:
-        imgs = [np.ascontiguousarray(data["img"]) for data in batch_data]
+    def _build_batch_sampler(self) -> ImageBatchSampler:
+        return ImageBatchSampler()
+
+    def _get_result_class(self) -> type:
+        return MLClassResult
+
+    def process(self, batch_data: List[Any]) -> Dict[str, List[Any]]:
+        batch_raw_imgs = self._data_reader(imgs=batch_data)
+        imgs = [np.ascontiguousarray(img) for img in batch_raw_imgs]
         ui_results = self._ui_model.batch_predict(imgs)
-        results: BatchData = []
-        for data, ui_result in zip(batch_data, ui_results):
-            ml_clas_result = self._create_ml_clas_result(data, ui_result)
-            results.append({"result": ml_clas_result})
-        return results
+
+        class_ids_list = []
+        scores_list = []
+        label_names_list = []
+        for ui_result in ui_results:
+            class_ids_list.append(ui_result.label_ids)
+            scores_list.append(np.around(ui_result.scores, decimals=5).tolist())
+            if self._label_list is not None:
+                label_names_list.append(
+                    [self._label_list[i] for i in ui_result.label_ids]
+                )
+
+        return {
+            "input_path": batch_data,
+            "input_img": batch_raw_imgs,
+            "class_ids": class_ids_list,
+            "scores": scores_list,
+            "label_names": label_names_list,
+        }
 
     def _get_label_list(self) -> Optional[List[str]]:
         pp_config = self.config["PostProcess"]
@@ -68,13 +89,3 @@ class MLClasPredictor(CVPredictor):
             raise RuntimeError("`MultiLabelThreshOutput` config not found")
         label_list = pp_config["MultiLabelThreshOutput"].get("label_list", None)
         return label_list
-
-    def _create_ml_clas_result(self, data: Data, ui_result: Any) -> MLClassResult:
-        dic = {
-            "input_path": data["input_path"],
-            "class_ids": ui_result.label_ids,
-            "scores": np.around(ui_result.scores, decimals=5).tolist(),
-        }
-        if self._label_list is not None:
-            dic["label_names"] = [self._label_list[i] for i in ui_result.label_ids]
-        return MLClassResult(dic)

libs/paddlex-hpi/src/paddlex_hpi/models/object_detection.py

@@ -15,13 +15,13 @@
 import os
 from typing import Any, Dict, List, Optional, Union
 
-import ultrainfer as ui
+import ultra_infer as ui
 import numpy as np
+from paddlex.inference.common.batch_sampler import ImageBatchSampler
 from paddlex.inference.results import DetResult
 from paddlex.modules.object_detection.model_list import MODELS
 from pydantic import BaseModel
 
-from paddlex_hpi._utils.typing import BatchData, Data
 from paddlex_hpi.models.base import CVPredictor, HPIParams
 
 
@@ -59,33 +59,30 @@ class DetPredictor(CVPredictor):
         )
         return model
 
-    def _predict(self, batch_data: BatchData) -> BatchData:
-        imgs = [np.ascontiguousarray(data["img"]) for data in batch_data]
-        ui_results = self._ui_model.batch_predict(imgs)
-        results: BatchData = []
-        for data, ui_result in zip(batch_data, ui_results):
-            det_result = self._create_det_result(data, ui_result)
-            results.append({"result": det_result})
-        return results
+    def _build_batch_sampler(self) -> ImageBatchSampler:
+        return ImageBatchSampler()
 
-    def _get_pp_params(self) -> _DetPPParams:
-        return _DetPPParams(
-            threshold=self.config["draw_threshold"],
-            label_list=self.config["label_list"],
-        )
+    def _get_result_class(self) -> type:
+        return DetResult
 
-    def _create_det_result(self, data: Data, ui_result: Any) -> DetResult:
-        inds = sorted(
-            range(len(ui_result.scores)), key=ui_result.scores.__getitem__, reverse=True
-        )
-        inds = [i for i in inds if ui_result.scores[i] > self._pp_params.threshold]
-        inds = [i for i in inds if ui_result.label_ids[i] > -1]
-        ids = [ui_result.label_ids[i] for i in inds]
-        scores = [ui_result.scores[i] for i in inds]
-        boxes = [ui_result.boxes[i] for i in inds]
-        dic = {
-            "input_path": data["input_path"],
-            "boxes": [
+    def process(self, batch_data: List[Any]) -> Dict[str, List[Any]]:
+        batch_raw_imgs = self._data_reader(imgs=batch_data)
+        imgs = [np.ascontiguousarray(img) for img in batch_raw_imgs]
+        ui_results = self._ui_model.batch_predict(imgs)
+
+        boxes_list = []
+        for ui_result in ui_results:
+            inds = sorted(
+                range(len(ui_result.scores)),
+                key=ui_result.scores.__getitem__,
+                reverse=True,
+            )
+            inds = [i for i in inds if ui_result.scores[i] > self._pp_params.threshold]
+            inds = [i for i in inds if ui_result.label_ids[i] > -1]
+            ids = [ui_result.label_ids[i] for i in inds]
+            scores = [ui_result.scores[i] for i in inds]
+            boxes = [ui_result.boxes[i] for i in inds]
+            boxes = [
                 {
                     "cls_id": id_,
                     "label": self._pp_params.label_list[id_],
@@ -93,6 +90,17 @@ class DetPredictor(CVPredictor):
                     "coordinate": box,
                 }
                 for id_, score, box in zip(ids, scores, boxes)
-            ],
+            ]
+            boxes_list.append(boxes)
+
+        return {
+            "input_path": batch_data,
+            "input_img": batch_raw_imgs,
+            "boxes": boxes_list,
         }
-        return DetResult(dic)
+
+    def _get_pp_params(self) -> _DetPPParams:
+        return _DetPPParams(
+            threshold=self.config["draw_threshold"],
+            label_list=self.config["label_list"],
+        )

libs/paddlex-hpi/src/paddlex_hpi/models/semantic_segmentation.py

@@ -12,14 +12,14 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Any, List
+from typing import Any, Dict, List
 
-import ultrainfer as ui
+import ultra_infer as ui
 import numpy as np
+from paddlex.inference.common.batch_sampler import ImageBatchSampler
 from paddlex.inference.results import SegResult
 from paddlex.modules.semantic_segmentation.model_list import MODELS
 
-from paddlex_hpi._utils.typing import BatchData, Data
 from paddlex_hpi.models.base import CVPredictor
 
 
@@ -37,20 +37,27 @@ class SegPredictor(CVPredictor):
         )
         return model
 
-    def _predict(self, batch_data: BatchData) -> BatchData:
-        imgs = [np.ascontiguousarray(data["img"]) for data in batch_data]
+    def _build_batch_sampler(self) -> ImageBatchSampler:
+        return ImageBatchSampler()
+
+    def _get_result_class(self) -> type:
+        return SegResult
+
+    def process(self, batch_data: List[Any]) -> Dict[str, List[Any]]:
+        batch_raw_imgs = self._data_reader(imgs=batch_data)
+        imgs = [np.ascontiguousarray(img) for img in batch_raw_imgs]
         ui_results = self._ui_model.batch_predict(imgs)
-        results: BatchData = []
-        for data, ui_result in zip(batch_data, ui_results):
-            seg_result = self._create_seg_result(data, ui_result)
-            results.append({"result": seg_result})
-        return results
-
-    def _create_seg_result(self, data: Data, ui_result: Any) -> SegResult:
-        pred = np.array(ui_result.label_map, dtype=np.int32).reshape(ui_result.shape)
-        pred = pred[np.newaxis]
-        dic = {
-            "input_path": data["input_path"],
-            "pred": pred,
+
+        batch_preds = []
+        for ui_result in ui_results:
+            pred = np.array(ui_result.label_map, dtype=np.int32).reshape(
+                ui_result.shape
+            )
+            pred = pred[np.newaxis]
+            batch_preds.append(pred)
+
+        return {
+            "input_path": batch_data,
+            "input_img": batch_raw_imgs,
+            "pred": batch_preds,
         }
-        return SegResult(dic)

libs/paddlex-hpi/src/paddlex_hpi/models/table_recognition.py

@@ -13,15 +13,15 @@
 # limitations under the License.
 
 import tempfile
-from typing import Any, List
+from typing import Any, Dict, List
 
-import ultrainfer as ui
+import ultra_infer as ui
 import numpy as np
+from paddlex.inference.common.batch_sampler import ImageBatchSampler
 from paddlex.inference.results import TableRecResult
 from paddlex.modules.table_recognition.model_list import MODELS
 
 from paddlex_hpi._utils.compat import get_compat_version
-from paddlex_hpi._utils.typing import BatchData, Data
 from paddlex_hpi.models.base import CVPredictor
 
 
@@ -50,19 +50,26 @@ class TablePredictor(CVPredictor):
             )
         return model
 
-    def _predict(self, batch_data: BatchData) -> BatchData:
-        imgs = [np.ascontiguousarray(data["img"]) for data in batch_data]
+    def _build_batch_sampler(self) -> ImageBatchSampler:
+        return ImageBatchSampler()
+
+    def _get_result_class(self) -> type:
+        return TableRecResult
+
+    def process(self, batch_data: List[Any]) -> Dict[str, List[Any]]:
+        batch_raw_imgs = self._data_reader(imgs=batch_data)
+        imgs = [np.ascontiguousarray(img) for img in batch_raw_imgs]
         ui_results = self._ui_model.batch_predict(imgs)
-        results: BatchData = []
-        for data, ui_result in zip(batch_data, ui_results):
-            table_result = self._create_table_result(data, ui_result)
-            results.append({"result": table_result})
-        return results
 
-    def _create_table_result(self, data: Data, ui_result: Any) -> TableRecResult:
-        dic = {
-            "input_path": data["input_path"],
-            "bbox": ui_result.table_boxes,
-            "structure": ui_result.table_structure,
+        bbox_list = []
+        structure_list = []
+        for ui_result in ui_results:
+            bbox_list.append(ui_result.table_boxes)
+            structure_list.append(ui_result.table_structure)
+
+        return {
+            "input_path": batch_data,
+            "input_img": batch_raw_imgs,
+            "bbox": bbox_list,
+            "structure": structure_list,
         }
-        return TableRecResult(dic)

libs/paddlex-hpi/src/paddlex_hpi/models/text_detection.py

@@ -15,13 +15,13 @@
 import os
 from typing import Any, Dict, List, Optional, Union
 
-import ultrainfer as ui
+import ultra_infer as ui
 import numpy as np
+from paddlex.inference.common.batch_sampler import ImageBatchSampler
 from paddlex.inference.results import TextDetResult
 from paddlex.modules.text_detection.model_list import CURVE_MODELS, MODELS
 
 from paddlex_hpi._utils.misc import parse_scale
-from paddlex_hpi._utils.typing import BatchData, Data
 from paddlex_hpi.models.base import CVPredictor, HPIParams
 
 
@@ -42,6 +42,12 @@ class TextDetPredictor(CVPredictor):
             hpi_params=hpi_params,
         )
 
+    def _build_batch_sampler(self) -> ImageBatchSampler:
+        return ImageBatchSampler()
+
+    def _get_result_class(self) -> type:
+        return TextDetResult
+
     # HACK
     @property
     def _is_curve_model(self) -> bool:
@@ -66,14 +72,27 @@ class TextDetPredictor(CVPredictor):
         self._config_ui_postprocessor(model)
         return model
 
-    def _predict(self, batch_data: BatchData) -> BatchData:
-        imgs = [np.ascontiguousarray(data["img"]) for data in batch_data]
+    def process(self, batch_data: List[Any]) -> Dict[str, List[Any]]:
+        batch_raw_imgs = self._data_reader(imgs=batch_data)
+        imgs = [np.ascontiguousarray(img) for img in batch_raw_imgs]
         ui_results = self._ui_model.batch_predict(imgs)
-        results: BatchData = []
-        for data, ui_result in zip(batch_data, ui_results):
-            text_det_result = self._create_text_det_result(data, ui_result)
-            results.append({"result": text_det_result})
-        return results
+
+        dt_polys_list = []
+        dt_scores_list = []
+        for ui_result in ui_results:
+            polys = [list(zip(*([iter(box)] * 2))) for box in ui_result.boxes]
+            dt_polys_list.append(polys)
+            # XXX: Currently, we cannot get scores from `ui_result`, so we
+            # temporarily use dummy scores here.
+            dummy_scores = [0.0 for _ in ui_result.boxes]
+            dt_scores_list.append(dummy_scores)
+
+        return {
+            "input_path": batch_data,
+            "input_img": batch_raw_imgs,
+            "dt_polys": dt_polys_list,
+            "dt_scores": dt_scores_list,
+        }
 
     def _config_ui_preprocessor(self, model: ui.vision.ocr.DBDetector) -> None:
         pp_config = self.config["PreProcess"]
@@ -153,15 +172,3 @@ class TextDetPredictor(CVPredictor):
                 postprocessor.det_db_box_type = "bbox"
             else:
                 postprocessor.det_db_box_type = "poly"
-
-    def _create_text_det_result(self, data: Data, ui_result: Any) -> TextDetResult:
-        polys = [list(zip(*([iter(box)] * 2))) for box in ui_result.boxes]
-        # XXX: Currently, we cannot get scores from `ui_result`, so we
-        # temporarily use dummy scores here.
-        dummy_scores = [0.0 for _ in ui_result.boxes]
-        dic = {
-            "input_path": data["input_path"],
-            "dt_polys": polys,
-            "dt_scores": dummy_scores,
-        }
-        return TextDetResult(dic)

libs/paddlex-hpi/src/paddlex_hpi/models/text_recognition.py

@@ -13,20 +13,26 @@
 # limitations under the License.
 
 import tempfile
-from typing import List
+from typing import Any, Dict, List
 
-import ultrainfer as ui
+import ultra_infer as ui
 import numpy as np
+from paddlex.inference.common.batch_sampler import ImageBatchSampler
 from paddlex.inference.results import TextRecResult
 from paddlex.modules.text_recognition.model_list import MODELS
 
-from paddlex_hpi._utils.typing import BatchData, Data
 from paddlex_hpi.models.base import CVPredictor
 
 
 class TextRecPredictor(CVPredictor):
     entities = MODELS
 
+    def _build_batch_sampler(self) -> ImageBatchSampler:
+        return ImageBatchSampler()
+
+    def _get_result_class(self) -> type:
+        return TextRecResult
+
     def _build_ui_model(self, option: ui.RuntimeOption) -> ui.vision.ocr.Recognizer:
         with tempfile.NamedTemporaryFile("w", encoding="utf-8", suffix=".txt") as f:
             pp_config = self.config["PostProcess"]
@@ -42,14 +48,20 @@ class TextRecPredictor(CVPredictor):
             self._config_ui_preprocessor(model)
         return model
 
-    def _predict(self, batch_data: BatchData) -> BatchData:
-        imgs = [np.ascontiguousarray(data["img"]) for data in batch_data]
-        ui_result = self._ui_model.batch_predict(imgs)
-        results: BatchData = []
-        for data, text, score in zip(batch_data, ui_result.text, ui_result.rec_scores):
-            text_rec_result = self._create_text_rec_result(data, text, score)
-            results.append({"result": text_rec_result})
-        return results
+    def process(self, batch_data: List[Any]) -> Dict[str, List[Any]]:
+        batch_raw_imgs = self._data_reader(imgs=batch_data)
+        imgs = [np.ascontiguousarray(img) for img in batch_raw_imgs]
+        ui_results = self._ui_model.batch_predict(imgs)
+
+        texts_list = ui_results.text
+        rec_score_list = ui_results.rec_scores
+
+        return {
+            "input_path": batch_data,
+            "input_img": batch_raw_imgs,
+            "rec_text": texts_list,
+            "rec_score": rec_score_list,
+        }
 
     def _config_ui_preprocessor(self, model: ui.vision.ocr.Recognizer) -> None:
         pp_config = self.config["PreProcess"]
@@ -74,13 +86,3 @@ class TextRecPredictor(CVPredictor):
                 raise RuntimeError(f"Unkown preprocessing operator: {op_name}")
         if not found_resize_op:
             raise RuntimeError("Could not find the config for `RecResizeImg`.")
-
-    def _create_text_rec_result(
-        self, data: Data, text: str, score: float
-    ) -> TextRecResult:
-        dic = {
-            "input_path": data["input_path"],
-            "rec_text": text,
-            "rec_score": score,
-        }
-        return TextRecResult(dic)

libs/paddlex-hpi/src/paddlex_hpi/models/ts_ad.py

@@ -14,7 +14,7 @@
 
 from typing import Any, List
 
-import ultrainfer as ui
+import ultra_infer as ui
 import pandas as pd
 from paddlex.inference.results import TSAdResult
 from paddlex.modules.ts_anomaly_detection.model_list import MODELS

libs/paddlex-hpi/src/paddlex_hpi/models/ts_cls.py

@@ -14,7 +14,7 @@
 
 from typing import Any, List
 
-import ultrainfer as ui
+import ultra_infer as ui
 import pandas as pd
 from paddlex.inference.results import TSClsResult
 from paddlex.modules.ts_classification.model_list import MODELS

libs/paddlex-hpi/src/paddlex_hpi/models/ts_fc.py

@@ -14,7 +14,7 @@
 
 from typing import Any, List
 
-import ultrainfer as ui
+import ultra_infer as ui
 import pandas as pd
 from paddlex.inference.results import TSFcResult
 from paddlex.modules.ts_forecast.model_list import MODELS