YOLOV8物体识别可以同时调用两个best.pt文件吗？

Question

我现在想利用yolov8做一个工地视觉识别系统，目前已经训练完成了一个识别安全帽的best.pt模型，还有一个火灾识别的best.pt模型。现在我想在实时视觉监测视频中同时检测安全帽与火灾，请问应该怎么办？

我知道可以重新通过打标签的方式把两个模型合成为一个模型，但是这样做的工作量实在太大了，请问有什么简便的方法吗？

Answer 1

阿里嘎多学长整理AIGC生成，因移动端显示问题导致当前答案未能完全显示，请使用PC端查看更加详细的解答过程

可以同时调用两个best.pt文件

YOLOV8支持多任务检测，理论上你可以同时调用两个best.pt文件来检测安全帽和火灾。具体步骤如下：

1. 加载模型: 使用 torch.hub.load() 或 torch.hub.load_pretrained() 加载两个best.pt模型。
2. 创建检测器: 使用 yolov8.load() 或 yolov8.load_from() 创建两个检测器实例。
3. 设置检测器参数: 设置两个检测器的参数，例如 conf_thres、iou_thres 等。
4. 检测视频: 使用 cv2.VideoCapture() 读取视频流，使用 yolov8.detect() 或 yolov8.run() 来检测安全帽和火灾。

核心代码

import torch
import cv2
from yolov8 import YOLOv8

# 加载模型
model1 = torch.hub.load('ultralytics/yolov8', 'yolov8n')
model2 = torch.hub.load('ultralytics/yolov8', 'yolov8n')

# 创建检测器
detector1 = YOLOv8(model1)
detector2 = YOLOv8(model2)

# 设置检测器参数
conf_thres = 0.5
iou_thres = 0.5

# 检测视频
cap = cv2.VideoCapture('video.mp4')
while True:
    ret, frame = cap.read()
    if not ret:
        break

    # 检测安全帽
    results1 = detector1(frame, conf_thres=conf_thres, iou_thres=iou_thres)
    # 检测火灾
    results2 = detector2(frame, conf_thres=conf_thres, iou_thres=iou_thres)

    # 绘制检测结果
    cv2.imshow('安全帽检测', results1.render())
    cv2.imshow('火灾检测', results2.render())

    if cv2.waitKey(1) & 0xFF == ord('q'):
        break

cap.release()
cv2.destroyAllWindows()

注意: 请确保两个模型的输入尺寸和检测器参数一致。