paddle_inference
import numpy as np
import cv2
import os
from ai.config.log import get_logger
from ai.env import alg_env
from io import BytesIO
引用 paddle inference 预测库
import paddle.inference as paddle_infer
logger = get_logger(name)
class Alg(object):
def init(self, alg_path=None):
logger.info(f"{'*' * 10} {alg_path}")
# paddle版本信息
print(paddle_infer.get_version())
base_dir = alg_env.get_py_base_dir()
if alg_path is None:
pass
alg_path = os.path.join(base_dir, "algs", "FloatingObjectDetection")
self.numtoen = {0: 'nothing', 1: 'plant',
2: 'garbage', 3: 'else'}
self.numtocn = {0: '无目标', 1: '植物类漂浮物',
2: '垃圾类漂浮物', 3: '其它类漂浮物'}
self.model_file = os.path.join(alg_path, "m/model.pdmodel")
self.params_file = os.path.join(alg_path, "m/model.pdiparams")
self.config_file = os.path.join(alg_path, "m/infer_cfg.yml")
# 模型初始化，返回是否成功
def init_model(self, resource):
# 根据给定路径加载模型
init_success = False
# 创建 config
config = paddle_infer.Config(self.model_file, self.params_file)
# 使用gpu预测，预测精度设置为FP16
config.enable_use_gpu(1000, 0)
config.exp_enable_use_gpu_fp16()
# 根据 config 创建 predictor
self.model = paddle_infer.create_predictor(config)
if self.model:
init_success = True
print('模型初始化成功')
return init_success
# 返回list[报警名，报警中文名，置信度，图片宽，图片高，*bbox]
def predict_img(self, img, task_id, img_w, img_h):
result_list = []
# 获取输入的名称
input_names = self.model.get_input_names()
# print("input_names:",input_names)
input_handle = self.model.get_input_handle(input_names[0])
# 从 CPU 获取数据，设置到 Tensor 内部
input_handle.copy_from_cpu(img)
# 执行predictor
self.model.run()
# 获取输出
output_names = self.model.get_output_names()
output_handle = self.model.get_output_handle(output_names[0])
output_data = output_handle.copy_to_cpu() # numpy.ndarray类型
# 结果整理
result_list = [output_names, output_handle, output_data]
# 结果保存
return result_list
def main():
alg = Alg()
if alg.init_model(""):
img_file = "Images/1.jpg"
content = None
with open(img_file, "rb") as f:
content = f.read()
raw_img = np.asarray(bytearray(BytesIO(content).read()), dtype="uint8")
img = cv2.imdecode(raw_img, cv2.IMREAD_COLOR)
img = cv2.resize(img, (640, 640))
# img = (img / 255.0 - 0.5) / 0.5
# img = (img / 255.0) / 1.0
img = img.transpose((2, 0, 1))
# 这里的astype要和初始化模型时声明的类型一致！
img = np.expand_dims(img, axis=0).astype('float16')
print(img.shape)
print(img)
data_result = alg.predict_img(img, 'xx0001', img.shape[1], img.shape[0])
print(f"result:{data_result}")
if name == "main":
main()
这是我的测试程序
是否意味着我是用GPU推理的时候一定要使用TRT