25 lines
1.1 KiB
Python
25 lines
1.1 KiB
Python
|
import paddle.inference as paddle_infer
|
||
|
|
import numpy as np
|
||
|
|
import paddle.vision.transforms as T
|
||
|
|
|
||
|
|
class Lane_model_infer:
|
||
|
|
def __init__(self, model_dir="./lane_model"):
|
||
|
|
# 初始化paddle推理
|
||
|
|
self.model_dir = model_dir
|
||
|
|
self.config = paddle_infer.Config(model_dir + "/model.pdmodel", model_dir + "/model.pdiparams")
|
||
|
|
self.config.enable_use_gpu(200, 0)
|
||
|
|
self.predictor = paddle_infer.create_predictor(self.config)
|
||
|
|
self.input_names = self.predictor.get_input_names()
|
||
|
|
self.input_handle = self.predictor.get_input_handle(self.input_names[0])
|
||
|
|
self.output_names = self.predictor.get_output_names()
|
||
|
|
self.output_handle = self.predictor.get_output_handle(self.output_names[0])
|
||
|
|
self.normalize_transform = T.Normalize(mean=[127.5], std=[127.5])
|
||
|
|
|
||
|
|
def infer(self,src) -> np.ndarray:
|
||
|
|
image = self.normalize_transform(src)
|
||
|
|
image = image.transpose(2, 0, 1)
|
||
|
|
image = np.expand_dims(image, axis=0)
|
||
|
|
self.input_handle.copy_from_cpu(image)
|
||
|
|
self.predictor.run()
|
||
|
|
results = self.output_handle.copy_to_cpu()[0]
|
||
|
|
return results
|