1.1 yolov5 使用 onnx模型

目前我们先试用yolov5官方训练好的模型，模型的后缀是.pt文件，我们可以把他转换为. onnx 方便python和c++ 来调用

首先需要有一个python的虚拟环境，需要下载 Anaconda3-2024.10-1-Windows-x86_64.exe

下载地址：

Anaconda3-2024.10-1-Windows-x86_64.exe - Information sharing - KodExplorer - Powered by KodExplorer

下载完成后，在anaconda的控制台运行

运行下来的3条命令，

conda install pytorch torchvision torchaudio cpuonly -c pytorch

conda install -c conda-forge opencv

conda install -c conda-forge onnx

基本就可以了

下来就是运行python代码

import torch

# 加载 YOLOv5x 模型
model = torch.hub.load('ultralytics/yolov5', 'yolov5x')

# 将模型设置为评估模式
model.eval()

# 输入的假数据（通常是一个与模型输入一致的张量）
dummy_input = torch.randn(1, 3, 640, 640)  # 1张图像，3个通道，640x640分辨率

# 导出为 ONNX 格式
onnx_file_path = "yolov5x.onnx"
torch.onnx.export(model,  # 模型
                  dummy_input,  # 假输入数据
                  onnx_file_path,  # 输出文件路径
                  opset_version=12,  # ONNX 版本
                  input_names=["input"],  # 输入的名字
                  output_names=["output"],  # 输出的名字
                  dynamic_axes={"input": {0: "batch_size"},  # 动态批次大小
                                "output": {0: "batch_size"}})
print(f"ONNX模型已保存为：{onnx_file_path}")

这样就初始化完成，创建了 onnx模型文件

下载在写运行另外一个python文件来做识别

import cv2
import numpy as np

# 加载 ONNX 模型
model_path = 'yolov5x.onnx'  # 替换为你的 ONNX 模型文件路径
net = cv2.dnn.readNetFromONNX(model_path)

# 读取输入图像
image_path = 'input_image.jpg'  # 替换为你要识别的图像路径
image = cv2.imread(image_path)

# 获取图像的高度和宽度
height, width = image.shape[:2]

# 将图像预处理为模型输入
blob = cv2.dnn.blobFromImage(image, 1 / 255.0, (640, 640), (0, 0, 0), swapRB=True, crop=False)
net.setInput(blob)

# 获取模型的输出
output = net.forward()

# 用于存储识别框信息的列表
class_ids = []
confidences = []
boxes = []

# 在 YOLO 模型的输出中进行解析
for out in output:
    for detection in out:
        scores = detection[5:]  # 假设前 5 个是框信息，后面的为各类别的置信度
        class_id = np.argmax(scores)
        confidence = scores[class_id]
        
        if confidence > 0.5:  # 设置置信度阈值
            center_x = int(detection[0] * width)  # 中心点 x
            center_y = int(detection[1] * height)  # 中心点 y
            w = int(detection[2] * width)  # 宽度
            h = int(detection[3] * height)  # 高度
            
            # 计算左上角坐标 (x1, y1) 和右下角坐标 (x2, y2)
            x1 = int(center_x - w / 2)
            y1 = int(center_y - h / 2)
          
            
            boxes.append([int(x1/1000), int(y1/1000), int(w/1000), int(h/1000)])  # 存储坐标和宽高
            confidences.append(float(confidence))
            class_ids.append(class_id)

# 输出检测到的目标数量
print(f"检测到 {len(boxes)} 个物体")

# 使用非极大值抑制（NMS）来去除重叠框
indices = cv2.dnn.NMSBoxes(boxes, confidences, 0.5, 0.4)

# 加载 COCO 类别名称（如果使用的是 YOLOv5 类似的模型）
with open('coco.names', 'r') as f:  # 替换为你的类别文件路径
    classes = f.read().strip().split('\n')

# 如果有检测框，打印出每个框的类别和置信度
if len(indices) > 0:
    print(f"总共有 {len(indices)} 个有效检测框（通过了NMS）")
else:
    print("没有检测到有效的物体")

# 绘制检测框并标注类别
for i in indices.flatten():
    x, y, w, h = boxes[i]
    label = str(classes[class_ids[i]])
    confidence = confidences[i]
    
    # 打印每个检测框的标签和置信度
    print(f"目标：{label}, 置信度：{confidence:.2f}, 坐标：({x}, {y}), 宽高：({w}, {h})")
    

    # 绘制矩形框
    color = (0, 255, 0)  # 绿色框
    cv2.rectangle(image, (x, y), (x + w, y + h), color, 2)
    
    # 添加标签和置信度
    #cv2.putText(image, f'{label} {confidence:.2f}', (x, y - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2)

# 保存带框的图像
output_image_path = 'output_image.jpg'  # 保存路径
cv2.imwrite(output_image_path, image)

生成识别后的文件

coco.names 的文件内容放在下面了 ，自己创建

person
bicycle
car
motorbike
aeroplane
bus
train
truck
boat
traffic light
fire hydrant
stop sign
parking meter
bench
bird
cat
dog
horse
sheep
cow
elephant
bear
zebra
giraffe
backpack
umbrella
handbag
tie
suitcase
frisbee
skis
snowboard
sports ball
kite
baseball bat
baseball glove
skateboard
surfboard
tennis racket
bottle
wine glass
cup
fork
knife
spoon
bowl
banana
apple
sandwich
orange
broccoli
carrot
hot dog
pizza
donut
cake
chair
sofa
pottedplant
bed
diningtable
toilet
tvmonitor
laptop
mouse
remote
keyboard
cell phone
microwave
oven
toaster
sink
refrigerator
book
clock
vase
scissors
teddy bear
hair drier
toothbrush

下一篇讲如何测试图片识别的效果