PaddleClas 车辆属性模型vehicle_attribute_model转onnx并部署

我变成了柴犬2025-05-18 20:39

下载paddlepaddle3.0.0

下载paddleclas 2.6.0

实验：运行前输入：

export FLAGS_enable_pir_api=0

python 复制代码

import paddleclas
model = paddleclas.PaddleClas(model_name="vehicle_attribute")
result = model.predict(input_data="/home/PaddleDetection-release-2.8.1/carimg/car.jpg")
print(next(result))

实验没问题后

替换成下载好的vehicle_attribute_model inference模型

python 复制代码

paddle2onnx --model_dir=./models/ResNet50_vd_infer/ \
--model_filename=inference.pdmodel \
--params_filename=inference.pdiparams \
--save_file=./models/ResNet50_vd_infer/inference.onnx \
--opset_version=10 \
--enable_onnx_checker=True

python 复制代码

import onnxruntime as ort
import numpy as np
import torchvision.transforms as transforms
from PIL import Image
session = ort.InferenceSession("/home/inference.onnx")
input_name = session.get_inputs()[0].name
output_name = session.get_outputs()[0].name
img = Image.open("/home/car2.jpg").convert('RGB')  
preprocess = transforms.Compose([
    transforms.Resize((192, 256)),  
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
input_tensor = preprocess(img)
input_batch = input_tensor.unsqueeze(0).numpy()  
result = session.run([output_name], {input_name: input_batch})
output = result[0][0]
color_map = {
    0: "yellow",
    1: "orange",
    2: "green",
    3: "gray",
    4: "red",
    5: "blue",
    6: "white",
    7: "golden",
    8: "brown",
    9:"black"
}
color_index = np.argmax(output[:10])  
color_prob = output[color_index]
output_dict = {
        "attributes": f"Color: ({color_map[color_index]}, prob: {color_prob})",
    "output": output.tolist(),
    "filename": "/home/PaddleDetection-release-2.8.1/carimg/car.jpg"
}

print(output_dict)