前言
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<