- 上一篇文章写了yolo的分类模型的训练,写篇文章基于yolo分类模型开发分类软件。
- 开发环境:pycharm,PySide6 6.6.1 ,PySide6-Addons 6.6.1,PySide6-Essentials 6.6.1,torch 2.3.1+cu121,torchaudio 2.3.1+cu121,torchvision 0.18.1+cu121,onnx 1.16.1,onnxruntime 1.17.3,opencv-contrib-python 4.10.0.82,opencv-python 4.10.0.82,opencv-python-headless 4.7.0.72
- 分类使用的数据集,halcon的pill分类demo的数据集
- 软件界面
- 核心代码
yolo推理代码
results = self.Model.predict(image_mat_)
if(len(results)>0):
###获取分类名
names_=results[0].names
probs_=results[0].probs
class_id_=probs_.top1
class_socre_=probs_.top1conf.item()
self.Classification=names_[class_id_]
self.Sorce=str(class_socre_)
onnx推理核心代码
target_image_height_ = input_onnx_.shape[2]###获取输入的目标维度
target_image_width_ = input_onnx_.shape[3]
# image_mat_=image_mat_.resize(target_image_width_,target_image_height_)###缩放图片
scale_image_mat_=cv2.resize(image_mat_,(target_image_width_,target_image_height_))
print("ImageMatShape:", scale_image_mat_.shape)
image_np_ = np.array(scale_image_mat_) ###图片转成np数组
print("ImageNpShape:", image_np_.shape)
image_np_ = image_np_.transpose(2, 0, 1) ##转成通道在前面的维度
print("ImageNpShape:", image_np_.shape)
resized_width = image_np_.shape[2] ##获取图片输入输出
resized_height = image_np_.shape[1]
image_np_ = image_np_.reshape(1, 3, target_image_height_, target_image_width_) ##添加一个新维度
print(image_np_[0, 0, 0, 0])
print("ImageNpShape:", image_np_.shape)
image_np_ = image_np_.astype(np.float32)
image_np_ = image_np_ / 255.0 ##数据归一化
print(image_np_[0, 0, 0, 0])
###获取输出点
outputs = self.OnnxModel.get_outputs()
output_onnx_ = outputs[0]
print("Name:", output_onnx_.name)
print("Type:", output_onnx_.type)
print("Shape:", output_onnx_.shape)
###运行推理
outputs = self.OnnxModel.run(["output0"], {"images": image_np_})
if(len(outputs)>0):
ng_sorce_=outputs[0][0][0]
ok_sorce_=outputs[0][0][1]
if(ng_sorce_>ok_sorce_):
self.Classification="ng"
self.Sorce=str(ng_sorce_)
else:
self.Classification="ok"
self.Sorce = str(ok_sorce_)
- 推理的效果显示