浏览代码

import uvdoc inference

zhangyubo0722 1 年之前
父节点
当前提交
1f2a25a4f0

+ 1 - 0
paddlex/inference/components/task_related/__init__.py

@@ -24,3 +24,4 @@ from .text_rec import OCRReisizeNormImg, CTCLabelDecode
 from .table_rec import TableLabelDecode
 from .table_rec import TableLabelDecode
 from .det import DetPostProcess, CropByBoxes
 from .det import DetPostProcess, CropByBoxes
 from .instance_seg import InstanceSegPostProcess
 from .instance_seg import InstanceSegPostProcess
+from .warp import DocTrPostProcess

+ 43 - 0
paddlex/inference/components/task_related/warp.py

@@ -0,0 +1,43 @@
+# 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 ..base import BaseComponent
+
+
+class DocTrPostProcess(BaseComponent):
+    """normalize image such as substract mean, divide std"""
+
+    INPUT_KEYS = ["pred"]
+    OUTPUT_KEYS = ["doctr_img"]
+    DEAULT_INPUTS = {"pred": "pred"}
+    DEAULT_OUTPUTS = {"doctr_img": "doctr_img"}
+
+    def __init__(self, scale=None, **kwargs):
+        super().__init__()
+        if isinstance(scale, str):
+            scale = np.float32(scale)
+        self.scale = np.float32(scale if scale is not None else 255.0)
+
+    def apply(self, pred):
+        im = pred[0]
+        assert isinstance(im, np.ndarray), "invalid input 'im' in DocTrPostProcess"
+
+        im = im.squeeze()
+        im = im.transpose(1, 2, 0)
+        im *= self.scale
+        im = im[:, :, ::-1]
+        im = im.astype("uint8", copy=False)
+        result = {"doctr_img": im}
+        return result

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

@@ -27,6 +27,7 @@ from .semantic_segmentation import SegPredictor
 from .general_recognition import ShiTuRecPredictor
 from .general_recognition import ShiTuRecPredictor
 from .ts_fc import TSFcPredictor
 from .ts_fc import TSFcPredictor
 from .ts_cls import TSClsPredictor
 from .ts_cls import TSClsPredictor
+from .image_unwarping import WarpPredictor
 
 
 
 
 def _create_hp_predictor(
 def _create_hp_predictor(

+ 51 - 0
paddlex/inference/predictors/image_unwarping.py

@@ -0,0 +1,51 @@
+# 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 ...modules.image_unwarping.model_list import MODELS
+from ..components import *
+from ..results import DocTrResult
+from ..utils.process_hook import batchable_method
+from .base import BasicPredictor
+
+
+class WarpPredictor(BasicPredictor):
+
+    entities = MODELS
+
+    def _check_args(self, kwargs):
+        assert set(kwargs.keys()).issubset(set(["batch_size"]))
+        return kwargs
+
+    def _build_components(self):
+        ops = {}
+        ops["ReadImage"] = ReadImage(
+            format="RGB", batch_size=self.kwargs.get("batch_size", 1)
+        )
+        ops["Normalize"] = Normalize(mean=0.0, std=1.0, scale=1.0 / 255)
+        ops["ToCHWImage"] = ToCHWImage()
+
+        predictor = ImagePredictor(
+            model_dir=self.model_dir,
+            model_prefix=self.MODEL_FILE_PREFIX,
+            option=self.pp_option,
+        )
+        ops["predictor"] = predictor
+
+        ops["postprocess"] = DocTrPostProcess()
+        return ops
+
+    @batchable_method
+    def _pack_res(self, single):
+        keys = ["img_path", "doctr_img"]
+        return DocTrResult({key: single[key] for key in keys})

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

@@ -22,3 +22,4 @@ from .det import DetResult
 from .seg import SegResult
 from .seg import SegResult
 from .instance_seg import InstanceSegResult
 from .instance_seg import InstanceSegResult
 from .ts import TSFcResult, TSClsResult
 from .ts import TSFcResult, TSClsResult
+from .warp import DocTrResult

+ 27 - 0
paddlex/inference/results/warp.py

@@ -0,0 +1,27 @@
+# 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 .base import BaseResult
+
+
+class DocTrResult(BaseResult):
+    def __init__(self, data):
+        super().__init__(data)
+        # We use opencv backend to save both numpy arrays
+        self._img_writer.set_backend("opencv")
+
+    def _get_res_img(self):
+        doctr_img = np.array(self["doctr_img"])
+        return doctr_img