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

本文链接：https://blog.csdn.net/qq_45640207/article/details/142992080

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

前言

对标的功能是【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 img
            img = 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:
            # 如果队列为空，则继续下一次循环（通常不会发生，因为设置了超时）
            continue
        except 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为空")
            # 如果没有结果，则继续下一次循环
            pass
        except Exception as e:
            print(f"处理推理结果时出现异常: {e}")

        # 显示结果
        # 显示结果
        frame_count += 1
        elapsed_time = time.time() - start_time
        fps = frame_count / elapsed_time
        cv2.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()