@@ -21,6 +21,8 @@ from .processors import LayoutAnalysisProcess
from .result import LayoutAnalysisResult
from .utils import STATIC_SHAPE_MODEL_LIST
+import os
+from ais_bench.infer.interface import InferSession
class LayoutAnalysisPredictor(DetPredictor):
@@ -61,6 +63,10 @@ class LayoutAnalysisPredictor(DetPredictor):
)
super().__init__(*args, **kwargs)
+ #init om
+ om_path = os.path.join(self.model_dir, f"PP-DocLayoutV3_linux_aarch64.om")
+ self.om_session = InferSession(0, om_path)
+
def _get_result_class(self):
return LayoutAnalysisResult
@@ -101,10 +107,13 @@ class LayoutAnalysisPredictor(DetPredictor):
batch_inputs = self.pre_ops[-1](datas)
# do infer
- batch_preds = self.infer(batch_inputs)
+ #batch_preds = self.infer(batch_inputs)
+ #om infer
+ batch_preds = self.om_session.infer(batch_inputs, mode='dymshape', custom_sizes=1000000000)
# process a batch of predictions into a list of single image result
preds_list = self._format_output(batch_preds)
+
# postprocess
boxes = self.post_op(
preds_list,
@@ -142,6 +142,7 @@ class DetPredictor(BasePredictor):
pre_ops.insert(1, self.build_resize(self.img_size, False, 2))
# build infer
+ '''
if self._use_static_model:
infer = self.create_static_infer()
else:
@@ -160,6 +161,8 @@ class DetPredictor(BasePredictor):
raise RuntimeError(
f"There is no dynamic graph implementation for model {repr(self.model_name)}."
)
+ '''
+ infer = None
# build postprocess op
post_op = self.build_postprocess()