【YOLO-v5】torch的使用

---

前言

torch调用已经训练好的模型进行识别操作

---

一、代码

初始版本

'''
#Author ：susocool
#Creattime:2024/9/18
#FileName:Hub_detect
#Description:加载一个自定义的 YOLOv5 模型导入

'''
import cv2
import torch
import sys
sys.path.insert(0, 'D:\\yolov5-master')  # 将 yolov5-master 目录添加到路径中
# 可以确保 Python 首先在这个目录中查找模块。
# 这对于解决模块名称冲突（例如，如果你有一个模块的副本）或者确保使用本地开发版本的模块非常有用。

# 加载本地YOLOv5模型
model_path = 'D:\\yolov5-master\\runs\\train\\exp4\\weights\\best.pt'
model = torch.hub.load('D:\\yolov5-master', 'custom', 
path=model_path, 
source='local')

# 初始化摄像头
cap = cv2.VideoCapture(0)
if not cap.isOpened():
    print("无法打开摄像头")
    exit()

while True:
    # 捕获摄像头的帧
    ret, frame = cap.read()
    if not ret:
        print("无法读取摄像头数据")
        break

    # 将图像转换为YOLOv5模型所需的格式
    img = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
    results = model(img)

    # 获取检测结果
    detections = results.pandas