当前位置：首页 > news >正文

【v5-Lite】模型导入使用-attempt_load

news 2025/6/18 16:30:40

前言
一、完整代码
二、运行结果
总结

前言

对标的功能是【YOLO-v5】torch的使用
但是在v5-lite里面是没有hubconf.py文件的（千万别跟我一样自作多情，转入然后走了很多弯路。缺这个少那个的）
正确的做法应该是attempt_load导入

一、完整代码

'''
#Author ：susocool
#Creattime:2024/10/16
#FileName:4-线程带异常处理
#Description: 1、帧数显示、加入异常处理2、线程处理-推理部分放在另一个线程中  
'''
import torch
import cv2
import numpy as np
from models.experimental import attempt_load
from utils.general import non_max_suppression, scale_coords
from utils.plots import plot_one_box
from queue import Queue, Empty
import threading
import time# 加载模型权重，指定CPU
try:model = attempt_load('D:\\YOLOv5-Lite-master\\v5lite-s.pt',map_location=torch.device('cpu'))model.eval()
except Exception as e:print(f"加载模型失败: {e}")exit()# 初始化摄像头
cap = cv2.VideoCapture(0)
if not cap.isOpened():print("无法打开摄像头")exit()# 创建一个队列来存储从摄像头捕获的帧
frame_queue = Queue()
# 创建一个队列来存储推理结果
result_queue = Queue()# 定义推理线程函数
def inference_thread(model, frame_queue, result_queue):while True:try:# 从队列中获取帧frame = frame_queue.get(timeout=1)  # 设置超时以避免阻塞if frame is None:  # 如果收到None，则退出线程break# 将图像转换为 YOLOv5 模型所需的格式global imgimg = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)img = torch.from_numpy(img).permute(2, 0, 1).float().div(255.0).unsqueeze(0)# 进行推理with torch.no_grad():pred = model(img)[0]# 非极大值抑制pred = non_max_suppression(pred, conf_thres=0.5, iou_thres=0.5)# 将结果放入结果队列result_queue.put(pred)except Empty:# 如果队列为空，则继续下一次循环（通常不会发生，因为设置了超时）continueexcept Exception as e:print(f"推理线程出现异常: {e}")# 启动推理线程
inference_thread_obj = threading.Thread(target=inference_thread, args=(model, frame_queue, result_queue))
inference_thread_obj.start()# 用于计算帧率的变量
frame_count = 0
start_time = time.time()try:while True:# 捕获摄像头的帧ret, frame = cap.read()if not ret:print("无法读取摄像头数据")break# 将帧放入队列以供推理线程处理try:frame_queue.put(frame)except Exception as e:print(f"将帧放入队列时出现异常: {e}")# 尝试从结果队列中获取推理结果try:pred = result_queue.get_nowait()  # 使用nowait以避免阻塞，如果没有结果则跳过绘制# 处理预测结果for det in pred:if det is not None and len(det):det[:, :4] = scale_coords(img.shape[2:], det[:, :4], frame.shape).round()for *xyxy, conf, cls in det:label = f'{model.names[int(cls)]} {conf:.2f}'label_onlyName = f'{model.names[int(cls)]}'plot_one_box(xyxy, frame, label=label, color=(0, 255, 0), line_thickness=3)print(label_onlyName)except Empty:# print("Empty为空")# 如果没有结果，则继续下一次循环passexcept Exception as e:print(f"处理推理结果时出现异常: {e}")# 显示结果# 显示结果frame_count += 1elapsed_time = time.time() - start_timefps = frame_count / elapsed_timecv2.putText( frame, f"FPS: {fps:.2f}", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 255, 0), 2, cv2.LINE_AA )cv2.imshow('v5lite识别结果', frame)if cv2.waitKey(1) & 0xFF == ord('q'):# 发送None到帧队列以通知推理线程停止frame_queue.put(None)break
finally:# 确保在退出前释放摄像头和关闭窗口cap.release()cv2.destroyAllWindows()# 等待推理线程结束inference_thread_obj.join()