YOLO5火焰识别的代码

以下是使用YOLOv5实现火焰检测的示例代码：

import cv2
import torch
from models.experimental import attempt_load
from utils.general import non_max_suppression, scale_coords
from utils.plots import plot_one_box

# 加载YOLOv5模型
model = attempt_load('yolov5s.pt', map_location=torch.device('cpu'))

# 设置阈值和IOU阈值
conf_thres = 0.25
iou_thres = 0.45

# 加载输入视频
cap = cv2.VideoCapture('fire_video.mp4')

while cap.isOpened():
    # 读取视频帧
    ret, frame = cap.read()

    if not ret:
        break

    # 将图像转换为RGB格式
    img = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)

    # 推理图像
    results = model(img)

    # 非最大抑制
    results = non_max_suppression(results, conf_thres, iou_thres)

    # 遍历检测结果
    for result in results:
        # 获取类别预测和坐标
        cls_pred, x1, y1, x2, y2, conf = result[0]

        # 缩放坐标
        scale_coords(img.shape[1:], [x1, y1, x2, y2], frame.shape).round()

        # 绘制检测框和类别标签
        label = f'{model.names[int(cls_pred)]} {conf:.2f}'
        plot_one_box([x1, y1, x2, y2], frame, label=label, color=(0, 255, 0))

    # 显示检测结果
    cv2.imshow('Fire Detection', frame)

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

# 释放资源
cap.release()
cv2.destroyAllWindows()

在此代码中，我们使用`yolov5s.pt`预训练模型进行火焰检测。输入视频名为`fire_video.mp4`。我们设置阈值为`conf_thres = 0.25`和IOU阈值为`iou_thres = 0.45`。在最终的检测结果中，我们使用`plot_one_box`函数绘制检测框和类别标签。