devbox
我家小花儿
这个屌丝很懒，什么也没留下！
关注作者
热门标签
热门文章
当前位置:   article > 正文

YoloV5 火焰与烟雾检测：深度学习在安全监控中的应用~附完整源码、超详细安装教程、烟雾检测系统、火灾检测数据集~火焰检测数据集~烟雾报警系统_烟火检测 yolo

作者：我家小花儿 | 2024-08-07 11:45:58

赞

踩

烟火检测 yolo

目录

  效果展示(完整源码请私信,并留下联系方式)

 数据集

 环境安装

环境安装：

YoloV5 简介

火焰与烟雾检测的挑战

YoloV5 的工作原理

实验与结果分析

应用前景与挑战

源码(完整源码请私信,并留下联系方式)

  效果展示(完整源码请私信,并留下联系方式)

基于YOLOv5的火焰与烟雾检测系统演示与介绍

 数据集

数据集准备了6744张已经标注好的数据集

 环境安装

runs文件夹中，存放训练和评估的结果图

环境安装：


请按照给定的python版本配置环境，否则可能会因依赖不兼容而出错，

在文件目录下cmd进入终端


(1)使用anaconda新建python3.10环境:
conda create -n env_rec python=3.10


(2)激活创建的环境:
conda activate env_rec


(3)使用pip安装所需的依赖，可通过requirements.txt:
pip install -r requirements.txt

在settings中找到project python interpreter 点击Add Interpreter

点击conda，在Use existing environment中选择刚才创建的虚拟环境 ，最后点击确定。如果conda Executable中路径没有，那就把anaconda3的路径添加上

随着人工智能和深度学习技术的快速发展，各行各业都在寻找将其应用于实际问题的新方法。其中，物体检测是一个热门的研究领域，而火灾和烟雾的实时检测尤为关键。在这篇博客中，我们将重点介绍 YoloV5 算法在火焰与烟雾检测中的应用及其技术细节。

YoloV5 简介

YoloV5 是一种基于深度学习的实时目标检测模型，由Ultralytics开发并于2020年发布。它建立在PyTorch框架上，通过一系列卷积神经网络层来检测图像中的多个物体。与早期的 Yolo 系列相比，YoloV5 在准确性和速度之间找到了良好的平衡，使其在实时应用中具有广泛的适用性。

火焰与烟雾检测的挑战

火灾和烟雾检测是安全监控中至关重要的任务，然而传统的方法往往依赖于规则和手工设计的特征，无法适应复杂多变的场景。深度学习模型，特别是像 YoloV5 这样的物体检测器，通过学习大量的图像数据和特征表示，能够更准确地识别火焰和烟雾，从而提高检测的精度和可靠性。

YoloV5 的工作原理

YoloV5 模型的核心是一个卷积神经网络，由一系列卷积层、池化层和全连接层组成。它能够以端到端的方式从原始图像中直接预测出目标的位置和类别。对于火焰和烟雾检测，YoloV5 需要训练数据集，其中包含了大量不同场景下的火灾和烟雾图像，以及相应的标签。

实验与结果分析

为了验证 YoloV5 在火焰与烟雾检测中的效果，我们使用了公开的数据集进行实验。在训练过程中，我们调整模型的超参数和数据增强策略，以提升检测的精度和鲁棒性。最终，我们评估了模型在测试集上的表现，并与其他经典方法进行了比较。实验结果显示，YoloV5 在火灾和烟雾检测中表现出色，具有较高的检测精度和较快的处理速度。

应用前景与挑战

尽管 YoloV5 在火灾与烟雾检测中取得了显著进展，但仍面临一些挑战。例如，不同场景下的光照条件、视角变化和目标尺度变化等问题仍然需要进一步研究和改进。此外，模型的实时性和稳定性对于安全监控系统的应用至关重要，需要在硬件支持和算法优化上持续努力。

综上所述，YoloV5 火焰与烟雾检测代表了深度学习在安全监控领域的重要应用。随着技术的进一步成熟和数据集的丰富，我们相信这些方法将在未来带来更多创新和改进，为社会安全和应急响应提供更多有力的支持。

通过本文的介绍，希望读者能对 YoloV5 在火灾与烟雾检测中的技术细节有所了解，并对其在实际应用中的潜力有所认识。未来，我们期待看到更多基于深度学习的物体检测技术在各个领域发挥重要作用，为人们的生活和安全保驾护航。

源码(完整源码请私信,并留下联系方式)

# -*- coding: UTF-8 -*-
"""
  @Author: mz
  @Date  : 2022/3/6 20:43
  @version V1.0
"""
import os
import random
import sys
import threading
import time

import cv2
import numpy
import torch
import torch.backends.cudnn as cudnn
from PyQt5.QtCore import *
from PyQt5.QtGui import *
from PyQt5.QtWidgets import *

from models.experimental import attempt_load
from utils.datasets import LoadImages, LoadStreams
from utils.general import check_img_size, non_max_suppression, scale_coords
from utils.plots import plot_one_box
from utils.torch_utils import select_device, time_synchronized

model_path = 'weights/best.pt'

# 添加一个关于界面
# 窗口主类
class MainWindow(QTabWidget):
    # 基本配置不动，然后只动第三个界面
    def __init__(self):
       # 初始化界面
       super().__init__()
       self.setWindowTitle('Yolov5火灾烟雾检测预警系统')
       self.resize(1200, 800)
       self.setWindowIcon(QIcon("./UI/xf.jpg"))
       # 图片读取进程
       self.output_size = 480
       self.img2predict = ""
       # 空字符串会自己进行选择，首选cuda
       self.device = ''
       # # 初始化视频读取线程
       self.vid_source = '0'  # 初始设置为摄像头
       # 检测视频的线程
       self.threading = None
       # 是否跳出当前循环的线程
       self.jump_threading: bool = False

       self.image_size = 640
       self.confidence = 0.25
       self.iou_threshold = 0.45
       # 指明模型加载的位置的设备
       self.model = self.model_load(weights=model_path,
                                    device=self.device)
       self.initUI()
       self.reset_vid()

    @torch.no_grad()
    def model_load(self,
                   weights="",  # model.pt path(s)
                   device='',  # cuda device, i.e. 0 or 0,1,2,3 or cpu
                   ):
       """
       模型初始化
       """
       device = self.device = select_device(device)
       # half = device.type != 'cpu'  # half precision only supported on CUDA
       half = device.type != 0
       # Load model
       model = attempt_load(weights, map_location=device)  # load FP32 model
       self.stride = int(model.stride.max())  # model stride
       self.image_size = check_img_size(self.image_size, s=self.stride)  # check img_size
       if half:
          model.half()  # to FP16
       # Run inference
       if device.type != 'cpu':
          print("Run inference")
          model(torch.zeros(1, 3, self.image_size, self.image_size).to(device).type_as(
             next(model.parameters())))  # run once
       print("模型加载完成!")
       return model

    def reset_vid(self):
       """
       界面重置事件
       """
       self.webcam_detection_btn.setEnabled(True)
       self.mp4_detection_btn.setEnabled(True)
       self.left_vid_img.setPixmap(QPixmap("./UI/up.jpeg"))
       self.vid_source = '0'
       self.disable_btn(self.det_img_button)
       self.disable_btn(self.vid_start_stop_btn)
       self.jump_threading = False

    def initUI(self):
       """
       界面初始化
       """
       # 图片检测子界面
       font_title = QFont('楷体', 16)
       font_main = QFont('楷体', 14)
       font_general = QFont('楷体', 10)
       # 图片识别界面, 两个按钮，上传图片和显示结果
       img_detection_widget = QWidget()
       img_detection_layout = QVBoxLayout()
       img_detection_title = QLabel("图片识别功能")
       img_detection_title.setFont(font_title)
       mid_img_widget = QWidget()
       mid_img_layout = QHBoxLayout()
       self.left_img = QLabel()
       self.right_img = QLabel()
       self.left_img.setPixmap(QPixmap("./UI/up.jpeg"))
       self.right_img.setPixmap(QPixmap("./UI/right.jpeg"))
       self.left_img.setAlignment(Qt.AlignCenter)
       self.right_img.setAlignment(Qt.AlignCenter)
       self.left_img.setMinimumSize(480, 480)
       self.left_img.setStyleSheet("QLabel{background-color: #f6f8fa;}")
       mid_img_layout.addWidget(self.left_img)
       self.right_img.setMinimumSize(480, 480)
       self.right_img.setStyleSheet("QLabel{background-color: #f6f8fa;}")
       mid_img_layout.addStretch(0)
       mid_img_layout.addWidget(self.right_img)
       mid_img_widget.setLayout(mid_img_layout)
       self.up_img_button = QPushButton("上传图片")
       self.det_img_button = QPushButton("开始检测")
       self.up_img_button.clicked.connect(self.upload_img)
       self.det_img_button.clicked.connect(self.detect_img)
       self.up_img_button.setFont(font_main)
       self.det_img_button.setFont(font_main)
       self.up_img_button.setStyleSheet("QPushButton{color:white}"
                                        "QPushButton:hover{background-color: rgb(2,110,180);}"
                                        "QPushButton{background-color:rgb(48,124,208)}"
                                        "QPushButton{border:2px}"
                                        "QPushButton{border-radius:5px}"
                                        "QPushButton{padding:5px 5px}"
                                        "QPushButton{margin:5px 5px}")
       self.det_img_button.setStyleSheet("QPushButton{color:white}"
                                         "QPushButton:hover{background-color: rgb(2,110,180);}"
                                         "QPushButton{background-color:rgb(48,124,208)}"
                                         "QPushButton{border:2px}"
                                         "QPushButton{border-radius:5px}"
                                         "QPushButton{padding:5px 5px}"
                                         "QPushButton{margin:5px 5px}")
       img_detection_layout.addWidget(img_detection_title, alignment=Qt.AlignCenter)
       img_detection_layout.addWidget(mid_img_widget, alignment=Qt.AlignCenter)
       img_detection_layout.addWidget(self.up_img_button)
       img_detection_layout.addWidget(self.det_img_button)
       img_detection_widget.setLayout(img_detection_layout)

       # 视频识别界面
       # 视频识别界面的逻辑比较简单，基本就从上到下的逻辑
       vid_detection_widget = QWidget()
       vid_detection_layout = QVBoxLayout()
       vid_title = QLabel("视频检测功能")
       vid_title.setFont(font_title)
       self.left_vid_img = QLabel()
       self.right_vid_img = QLabel()
       self.left_vid_img.setPixmap(QPixmap("./UI/up.jpeg"))
       self.right_vid_img.setPixmap(QPixmap("./UI/right.jpeg"))
       self.left_vid_img.setAlignment(Qt.AlignCenter)
       self.left_vid_img.setMinimumSize(480, 480)
       self.left_vid_img.setStyleSheet("QLabel{background-color: #f6f8fa;}")
       self.right_vid_img.setAlignment(Qt.AlignCenter)
       self.right_vid_img.setMinimumSize(480, 480)
       self.right_vid_img.setStyleSheet("QLabel{background-color: #f6f8fa;}")
       mid_img_widget = QWidget()
       mid_img_layout = QHBoxLayout()
       mid_img_layout.addWidget(self.left_vid_img)
       mid_img_layout.addStretch(0)
       mid_img_layout.addWidget(self.right_vid_img)
       mid_img_widget.setLayout(mid_img_layout)
       self.webcam_detection_btn = QPushButton("摄像头实时监测")
       self.mp4_detection_btn = QPushButton("视频文件检测")
       self.vid_start_stop_btn = QPushButton("启动/停止检测")
       self.webcam_detection_btn.setFont(font_main)
       self.mp4_detection_btn.setFont(font_main)
       self.vid_start_stop_btn.setFont(font_main)
       self.webcam_detection_btn.setStyleSheet("QPushButton{color:white}"
                                               "QPushButton:hover{background-color: rgb(2,110,180);}"
                                               "QPushButton{background-color:rgb(48,124,208)}"
                                               "QPushButton{border:2px}"
                                               "QPushButton{border-radius:5px}"
                                               "QPushButton{padding:5px 5px}"
                                               "QPushButton{margin:5px 5px}")
       self.mp4_detection_btn.setStyleSheet("QPushButton{color:white}"
                                            "QPushButton:hover{background-color: rgb(2,110,180);}"
                                            "QPushButton{background-color:rgb(48,124,208)}"
                                            "QPushButton{border:2px}"
                                            "QPushButton{border-radius:5px}"
                                            "QPushButton{padding:5px 5px}"
                                            "QPushButton{margin:5px 5px}")
       self.vid_start_stop_btn.setStyleSheet("QPushButton{color:white}"
                                             "QPushButton:hover{background-color: rgb(2,110,180);}"
                                             "QPushButton{background-color:rgb(48,124,208)}"
                                             "QPushButton{border:2px}"
                                             "QPushButton{border-radius:5px}"
                                             "QPushButton{padding:5px 5px}"
                                             "QPushButton{margin:5px 5px}")
       self.webcam_detection_btn.clicked.connect(self.open_cam)
       self.mp4_detection_btn.clicked.connect(self.open_mp4)
       self.vid_start_stop_btn.clicked.connect(self.start_or_stop)

       # 添加fps显示
       fps_container = QWidget()
       fps_container.setStyleSheet("QWidget{background-color: #f6f8fa;}")
       fps_container_layout = QHBoxLayout()
       fps_container.setLayout(fps_container_layout)
       # 左容器
       fps_left_container = QWidget()
       fps_left_container.setStyleSheet("QWidget{background-color: #f6f8fa;}")
       fps_left_container_layout = QHBoxLayout()
       fps_left_container.setLayout(fps_left_container_layout)

       # 右容器
       fps_right_container = QWidget()
       fps_right_container.setStyleSheet("QWidget{background-color: #f6f8fa;}")
       fps_right_container_layout = QHBoxLayout()
       fps_right_container.setLayout(fps_right_container_layout)

       # 将左容器和右容器添加到fps_container_layout中
       fps_container_layout.addWidget(fps_left_container)
       fps_container_layout.addStretch(0)
       fps_container_layout.addWidget(fps_right_container)

       # 左容器中添加fps显示
       raw_fps_label = QLabel("原始帧率:")
       raw_fps_label.setFont(font_general)
       raw_fps_label.setAlignment(Qt.AlignLeft)
       raw_fps_label.setStyleSheet("QLabel{margin-left:80px}")
       self.raw_fps_value = QLabel("0")
       self.raw_fps_value.setFont(font_general)
       self.raw_fps_value.setAlignment(Qt.AlignLeft)
       fps_left_container_layout.addWidget(raw_fps_label)
       fps_left_container_layout.addWidget(self.raw_fps_value)

       # 右容器中添加fps显示
       detect_fps_label = QLabel("检测帧率:")
       detect_fps_label.setFont(font_general)
       detect_fps_label.setAlignment(Qt.AlignRight)
       self.detect_fps_value = QLabel("0")
       self.detect_fps_value.setFont(font_general)
       self.detect_fps_value.setAlignment(Qt.AlignRight)
       self.detect_fps_value.setStyleSheet("QLabel{margin-right:96px}")
       fps_right_container_layout.addWidget(detect_fps_label)
       fps_right_container_layout.addWidget(self.detect_fps_value)

       # 添加组件到布局上
       vid_detection_layout.addWidget(vid_title, alignment=Qt.AlignCenter)
       vid_detection_layout.addWidget(fps_container)
       vid_detection_layout.addWidget(mid_img_widget, alignment=Qt.AlignCenter)
       vid_detection_layout.addWidget(self.webcam_detection_btn)
       vid_detection_layout.addWidget(self.mp4_detection_btn)
       vid_detection_layout.addWidget(self.vid_start_stop_btn)
       vid_detection_widget.setLayout(vid_detection_layout)

       # 关于界面
       about_widget = QWidget()
       about_layout = QVBoxLayout()
       about_title = QLabel('欢迎使用目标检测系统\n\n 可以进行知识交流\n\n wx:sybh0117')  # 修改欢迎词语
       about_title.setFont(QFont('楷体', 18))
       about_title.setAlignment(Qt.AlignCenter)
       about_img = QLabel()
       about_img.setPixmap(QPixmap('./UI/qq.png'))
       about_img.setAlignment(Qt.AlignCenter)


       label_super = QLabel()  # 更换作者信息
       label_super.setText("<a href='https://blog.csdn.net/m0_68036862?type=blog'>或者你可以在这里找到我-->mz</a>")
       label_super.setFont(QFont('楷体', 16))
       label_super.setOpenExternalLinks(True)
       # label_super.setOpenExternalLinks(True)
       label_super.setAlignment(Qt.AlignRight)
       about_layout.addWidget(about_title)
       about_layout.addStretch()
       about_layout.addWidget(about_img)
       about_layout.addStretch()
       about_layout.addWidget(label_super)
       about_widget.setLayout(about_layout)

       self.addTab(img_detection_widget, '图片检测')
       self.addTab(vid_detection_widget, '视频检测')
       self.addTab(about_widget, '联系我')
       self.setTabIcon(0, QIcon('./UI/lufei.png'))
       self.setTabIcon(1, QIcon('./UI/lufei.png'))

    def disable_btn(self, pushButton: QPushButton):
       pushButton.setDisabled(True)
       pushButton.setStyleSheet("QPushButton{background-color: rgb(2,110,180);}")

    def enable_btn(self, pushButton: QPushButton):
       pushButton.setEnabled(True)
       pushButton.setStyleSheet(
          "QPushButton{background-color: rgb(48,124,208);}"
          "QPushButton{color:white}"
       )

    def detect(self, source: str, left_img: QLabel, right_img: QLabel):
       """
       @param source: file/dir/URL/glob, 0 for webcam
       @param left_img: 将左侧QLabel对象传入，用于显示图片
       @param right_img: 将右侧QLabel对象传入，用于显示图片
       """
       model = self.model
       img_size = [self.image_size, self.image_size]  # inference size (pixels)
       conf_threshold = self.confidence  # confidence threshold
       iou_threshold = self.iou_threshold  # NMS IOU threshold
       device = self.device  # cuda device, i.e. 0 or 0,1,2,3 or cpu
       classes = None # filter by class: --class 0, or --class 0 2 3
       agnostic_nms = False  # class-agnostic NMS
       augment = False  # augmented inference

       half = device.type != 'cpu'  # half precision only supported on CUDA

       if source == "":
          self.disable_btn(self.det_img_button)
          QMessageBox.warning(self, "请上传", "请先上传视频或图片再进行检测")
       else:
          source = str(source)
          webcam = source.isnumeric()

          # Set Dataloader
          if webcam:
             cudnn.benchmark = True  # set True to speed up constant image size inference
             dataset = LoadStreams(source, img_size=img_size, stride=self.stride)
          else:
             dataset = LoadImages(source, img_size=img_size, stride=self.stride)
          # Get names and colors
          names = model.module.names if hasattr(model, 'module') else model.names
          colors = [[random.randint(0, 255) for _ in range(3)] for _ in names]

          # 用来记录处理的图片数量
          count = 0
          # 计算帧率开始时间
          fps_start_time = time.time()
          for path, img, im0s, vid_cap in dataset:
             # 直接跳出for，结束线程
             if self.jump_threading:
                # 清除状态
                self.jump_threading = False
                break
             count += 1
             img = torch.from_numpy(img).to(device)
             img = img.half() if half else img.float()  # uint8 to fp16/32
             img /= 255.0  # 0 - 255 to 0.0 - 1.0
             if img.ndimension() == 3:
                img = img.unsqueeze(0)

             # Inference
             t1 = time_synchronized()
             pred = model(img, augment=augment)[0]

             # Apply NMS
             pred = non_max_suppression(pred, conf_threshold, iou_threshold, classes=classes, agnostic=agnostic_nms)
             t2 = time_synchronized()

             # Process detections
             for i, det in enumerate(pred):  # detections per image
                if webcam:  # batch_size >= 1
                   s, im0 = 'detect : ', im0s[i].copy()
                else:
                   s, im0 = 'detect : ', im0s.copy()

                # s += '%gx%g ' % img.shape[2:]  # print string
                if len(det):
                   # Rescale boxes from img_size to im0 size
                   det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()

                   # Print results
                   for c in det[:, -1].unique():
                      n = (det[:, -1] == c).sum()  # detections per class
                      s += f"{n} {names[int(c)]}{'s' * (n > 1)}, "  # add to string

                   # Write results
                   for *xyxy, conf, cls in reversed(det):
                      label = f'{names[int(cls)]} {conf:.2f}'
                      plot_one_box(xyxy, im0, label=label, color=colors[int(cls)], line_thickness=3)
                      if names[int(cls)]=="fire" or names[int(cls)]=="smoke":
                         im0 = cv2.putText(im0, "Warning", (50, 110), cv2.FONT_HERSHEY_SIMPLEX,
                                       1, (0, 0, 255), 2, cv2.LINE_AA)


                if webcam or vid_cap is not None:
                   if webcam:  # batch_size >= 1
                      img = im0s[i]
                   else:
                      img = im0s
                   img = self.resize_img(img)
                   img = QImage(img.data, img.shape[1], img.shape[0], img.shape[2] * img.shape[1],
                                QImage.Format_RGB888)
                   left_img.setPixmap(QPixmap.fromImage(img))
                   # 计算一次帧率
                   if count % 10 == 0:
                      fps = int(10 / (time.time() - fps_start_time))
                      self.detect_fps_value.setText(str(fps))
                      fps_start_time = time.time()
                # 应该调整一下图片的大小
                # 时间显示
                timenumber = time.strftime('%Y/%m/%d/-%H:%M:%S', time.localtime(time.time()))
                im0 = cv2.putText(im0, timenumber, (50, 50), cv2.FONT_HERSHEY_SIMPLEX,
                              1, (0, 255, 0), 2, cv2.LINE_AA)
                im0 = cv2.putText(im0, s, (50, 80), cv2.FONT_HERSHEY_SIMPLEX,
                              1, (255, 0, 0), 2, cv2.LINE_AA)
                img = self.resize_img(im0)
                img = QImage(img.data, img.shape[1], img.shape[0], img.shape[2] * img.shape[1],
                             QImage.Format_RGB888)
                right_img.setPixmap(QPixmap.fromImage(img))

                # Print time (inference + NMS)
                print(f'{s}Done. ({t2 - t1:.3f}s)')

          # 使用完摄像头释放资源
          if webcam:
             for cap in dataset.caps:
                cap.release()
          else:
             dataset.cap and dataset.cap.release()

    def resize_img(self, img):
       """
       调整图片大小，方便用来显示
       @param img: 需要调整的图片
       """
       resize_scale = min(self.output_size / img.shape[0], self.output_size / img.shape[1])
       img = cv2.resize(img, (0, 0), fx=resize_scale, fy=resize_scale)
       img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
       return img

    def upload_img(self):
       """
       上传图片
       """
       # # 选择录像文件进行读取
       # fileName, fileType = QFileDialog.getOpenFileName(self, 'Choose file', '', '*.jpg *.png *.tif *.jpeg')
       # if fileName:
       #  self.img2predict = fileName
       #  # 将上传照片和进行检测做成互斥的
       #  self.enable_btn(self.det_img_button)
       #  self.disable_btn(self.up_img_button)
       #  # 进行左侧原图展示
       #  img = cv2.imread(fileName)
       #  # 应该调整一下图片的大小
       #  img = self.resize_img(img)
       #  img = QImage(img.data, img.shape[1], img.shape[0], img.shape[2] * img.shape[1], QImage.Format_RGB888)
       #  self.left_img.setPixmap(QPixmap.fromImage(img))
       #  # 上传图片之后右侧的图片重置
       #  self.right_img.setPixmap(QPixmap("./UI/right.jpeg"))
       fileName, fileType = QFileDialog.getOpenFileName(self, 'Choose file', '', '*.jpg *.png *.tif *.jpeg')

       if fileName:
          # 检查文件是否存在
          if not os.path.exists(fileName):
             print("File does not exist:", fileName)
             return

          self.img2predict = fileName
          # 设置按钮状态
          self.enable_btn(self.det_img_button)
          self.disable_btn(self.up_img_button)
          # 读取图像文件
          img = cv2.imread(fileName)


          if img is None:
             print(fileName, fileType)
             print("Error: Could not read image. Check file format or integrity.")
             return

          # 调整图像大小
          img = self.resize_img(img)
          # 转换为 QImage
          height, width, channel = img.shape
          bytesPerLine = channel * width
          qImg = QImage(img.data, width, height, bytesPerLine, QImage.Format_RGB888)
          # 设置左侧图像展示
          self.left_img.setPixmap(QPixmap.fromImage(qImg))
          # 重置右侧图像展示
          self.right_img.setPixmap(QPixmap("./UI/right.jpeg"))
       else:
          print("No file selected.")
    def detect_img(self):
       """
       检测图片
       """
       # 重置跳出线程状态，防止其他位置使用的影响
       self.jump_threading = False
       self.detect(self.img2predict, self.left_img, self.right_img)
       # 将上传照片和进行检测做成互斥的
       self.enable_btn(self.up_img_button)
       self.disable_btn(self.det_img_button)

    def open_mp4(self):
       """
       开启视频文件检测事件
       """
       print("开启视频文件检测")
       fileName, fileType = QFileDialog.getOpenFileName(self, 'Choose file', '', '*.mp4 *.avi')
       if fileName:
          self.disable_btn(self.webcam_detection_btn)
          self.disable_btn(self.mp4_detection_btn)
          self.enable_btn(self.vid_start_stop_btn)
          # 生成读取视频对象
          print(fileName)
          cap = cv2.VideoCapture(fileName)
          # 获取视频的帧率
          fps = cap.get(cv2.CAP_PROP_FPS)
          # 显示原始视频帧率
          self.raw_fps_value.setText(str(fps))
          if cap.isOpened():
             # 读取一帧用来提前左侧展示
             ret, raw_img = cap.read()
             cap.release()
          else:
             QMessageBox.warning(self, "需要重新上传", "请重新选择视频文件")
             self.disable_btn(self.vid_start_stop_btn)
             self.enable_btn(self.webcam_detection_btn)
             self.enable_btn(self.mp4_detection_btn)
             return
          # 应该调整一下图片的大小
          img = self.resize_img(numpy.array(raw_img))
          img = QImage(img.data, img.shape[1], img.shape[0], img.shape[2] * img.shape[1], QImage.Format_RGB888)
          self.left_vid_img.setPixmap(QPixmap.fromImage(img))
          # 上传图片之后右侧的图片重置
          self.right_vid_img.setPixmap(QPixmap("./UI/right.jpeg"))
          self.vid_source = fileName
          self.jump_threading = False

    def open_cam(self):
       """
       打开摄像头事件
       """
       print("打开摄像头")
       self.disable_btn(self.webcam_detection_btn)
       self.disable_btn(self.mp4_detection_btn)
       self.enable_btn(self.vid_start_stop_btn)
       self.vid_source = "0"
       self.jump_threading = False
       # 生成读取视频对象
       cap = cv2.VideoCapture(0)
       # 获取视频的帧率
       fps = cap.get(cv2.CAP_PROP_FPS)
       # 显示原始视频帧率
       self.raw_fps_value.setText(str(fps))
       if cap.isOpened():
          # 读取一帧用来提前左侧展示
          ret, raw_img = cap.read()
          cap.release()
       else:
          QMessageBox.warning(self, "需要重新上传", "请重新选择视频文件")
          self.disable_btn(self.vid_start_stop_btn)
          self.enable_btn(self.webcam_detection_btn)
          self.enable_btn(self.mp4_detection_btn)
          return
       # 应该调整一下图片的大小
       img = self.resize_img(numpy.array(raw_img))
       img = QImage(img.data, img.shape[1], img.shape[0], img.shape[2] * img.shape[1], QImage.Format_RGB888)
       self.left_vid_img.setPixmap(QPixmap.fromImage(img))
       # 上传图片之后右侧的图片重置
       self.right_vid_img.setPixmap(QPixmap("./UI/right.jpeg"))

    def start_or_stop(self):
       """
       启动或者停止事件
       """
       print("启动或者停止")
       if self.threading is None:
          # 创造并启动一个检测视频线程
          self.jump_threading = False
          self.threading = threading.Thread(target=self.detect_vid)
          self.threading.start()
          self.disable_btn(self.webcam_detection_btn)
          self.disable_btn(self.mp4_detection_btn)
       else:
          # 停止当前线程
          # 线程属性置空，恢复状态
          self.threading = None
          self.jump_threading = True
          self.enable_btn(self.webcam_detection_btn)
          self.enable_btn(self.mp4_detection_btn)

    def detect_vid(self):
       """
       视频检测
       视频和摄像头的主函数是一样的，不过是传入的source不同罢了
       """
       print("视频开始检测")
       self.detect(self.vid_source, self.left_vid_img, self.right_vid_img)
       print("视频检测结束")
       # 执行完进程，刷新一下和进程有关的状态，只有self.threading是None，
       # 才能说明是正常结束的线程，需要被刷新状态
       if self.threading is not None:
          self.start_or_stop()

    def closeEvent(self, event):
       """
       界面关闭事件
       """
       reply = QMessageBox.question(
          self,
          'quit',
          "Are you sure?",
          QMessageBox.Yes | QMessageBox.No,
          QMessageBox.No
       )
       if reply == QMessageBox.Yes:
          self.jump_threading = True
          self.close()
          event.accept()
       else:
          event.ignore()


if __name__ == "__main__":
    app = QApplication(sys.argv)
    mainWindow = MainWindow()
    mainWindow.show()
    sys.exit(app.exec_())

训练代码(train.py)

  1. import argparse
  2. import logging
  3. import math
  4. import os
  5. import random
  6. import time
  7. from copy import deepcopy
  8. from pathlib import Path
  9. from threading import Thread
  10.  
  11. import numpy as np
  12. import torch.distributed as dist
  13. import torch.nn as nn
  14. import torch.nn.functional as F
  15. import torch.optim as optim
  16. import torch.optim.lr_scheduler as lr_scheduler
  17. import torch.utils.data
  18. import yaml
  19. from torch.cuda import amp
  20. from torch.nn.parallel import DistributedDataParallel as DDP
  21. from torch.utils.tensorboard import SummaryWriter
  22. from tqdm import tqdm
  23.  
  24. import test  # import test.py to get mAP after each epoch
  25. from models.experimental import attempt_load
  26. from models.yolo import Model
  27. from utils.autoanchor import check_anchors
  28. from utils.datasets import create_dataloader
  29. from utils.general import labels_to_class_weights, increment_path, labels_to_image_weights, init_seeds, \
  30.     fitness, strip_optimizer, get_latest_run, check_dataset, check_file, check_git_status, check_img_size, \
  31.     check_requirements, print_mutation, set_logging, one_cycle, colorstr
  32. from utils.google_utils import attempt_download
  33. from utils.loss import ComputeLoss
  34. from utils.plots import plot_images, plot_labels, plot_results, plot_evolution
  35. from utils.torch_utils import ModelEMA, select_device, intersect_dicts, torch_distributed_zero_first, is_parallel
  36. from utils.wandb_logging.wandb_utils import WandbLogger, check_wandb_resume
  37.  
  38. logger = logging.getLogger(__name__)
  39.  
  40.  
  41. def train(hyp, opt, device, tb_writer=None):
  42.     logger.info(colorstr('hyperparameters: ') + ', '.join(f'{k}={v}' for k, v in hyp.items()))
  43.     save_dir, epochs, batch_size, total_batch_size, weights, rank = \
  44.         Path(opt.save_dir), opt.epochs, opt.batch_size, opt.total_batch_size, opt.weights, opt.global_rank
  45.  
  46.     # Directories
  47.     wdir = save_dir / 'weights'
  48.     wdir.mkdir(parents=True, exist_ok=True)  # make dir
  49.     last = wdir / 'last.pt'
  50.     best = wdir / 'best.pt'
  51.     results_file = save_dir / 'results.txt'
  52.  
  53.     # Save run settings
  54.     with open(save_dir / 'hyp.yaml', 'w') as f:
  55.         yaml.dump(hyp, f, sort_keys=False)
  56.     with open(save_dir / 'opt.yaml', 'w') as f:
  57.         yaml.dump(vars(opt), f, sort_keys=False)
  58.  
  59.     # Configure
  60.     plots = not opt.evolve  # create plots
  61.     cuda = device.type != 'cpu'
  62.     init_seeds(2 + rank)
  63.     with open(opt.data) as f:
  64.         data_dict = yaml.load(f, Loader=yaml.SafeLoader)  # data dict
  65.     is_coco = opt.data.endswith('coco.yaml')
  66.  
  67.     # Logging- Doing this before checking the dataset. Might update data_dict
  68.     loggers = {'wandb': None}  # loggers dict
  69.     if rank in [-1, 0]:
  70.         opt.hyp = hyp  # add hyperparameters
  71.         run_id = torch.load(weights).get('wandb_id') if weights.endswith('.pt') and os.path.isfile(weights) else None
  72.         wandb_logger = WandbLogger(opt, Path(opt.save_dir).stem, run_id, data_dict)
  73.         loggers['wandb'] = wandb_logger.wandb
  74.         data_dict = wandb_logger.data_dict
  75.         if wandb_logger.wandb:
  76.             weights, epochs, hyp = opt.weights, opt.epochs, opt.hyp  # WandbLogger might update weights, epochs if resuming
  77.  
  78.     nc = 1 if opt.single_cls else int(data_dict['nc'])  # number of classes
  79.     names = ['item'] if opt.single_cls and len(data_dict['names']) != 1 else data_dict['names']  # class names
  80.     assert len(names) == nc, '%g names found for nc=%g dataset in %s' % (len(names), nc, opt.data)  # check
  81.  
  82.     # Model
  83.     pretrained = weights.endswith('.pt')
  84.     if pretrained:
  85.         with torch_distributed_zero_first(rank):
  86.             attempt_download(weights)  # download if not found locally
  87.         ckpt = torch.load(weights, map_location=device)  # load checkpoint
  88.         model = Model(opt.cfg or ckpt['model'].yaml, ch=3, nc=nc, anchors=hyp.get('anchors')).to(device)  # create
  89.         exclude = ['anchor'] if (opt.cfg or hyp.get('anchors')) and not opt.resume else []  # exclude keys
  90.         state_dict = ckpt['model'].float().state_dict()  # to FP32
  91.         state_dict = intersect_dicts(state_dict, model.state_dict(), exclude=exclude)  # intersect
  92.         model.load_state_dict(state_dict, strict=False)  # load
  93.         logger.info('Transferred %g/%g items from %s' % (len(state_dict), len(model.state_dict()), weights))  # report
  94.     else:
  95.         model = Model(opt.cfg, ch=3, nc=nc, anchors=hyp.get('anchors')).to(device)  # create
  96.     with torch_distributed_zero_first(rank):
  97.         check_dataset(data_dict)  # check
  98.     train_path = data_dict['train']
  99.     test_path = data_dict['val']
  100.  
  101.     # Freeze
  102.     freeze = []  # parameter names to freeze (full or partial)
  103.     for k, v in model.named_parameters():
  104.         v.requires_grad = True  # train all layers
  105.         if any(x in k for x in freeze):
  106.             print('freezing %s' % k)
  107.             v.requires_grad = False
  108.  
  109.     # Optimizer
  110.     nbs = 64  # nominal batch size
  111.     accumulate = max(round(nbs / total_batch_size), 1)  # accumulate loss before optimizing
  112.     hyp['weight_decay'] *= total_batch_size * accumulate / nbs  # scale weight_decay
  113.     logger.info(f"Scaled weight_decay = {hyp['weight_decay']}")
  114.  
  115.     pg0, pg1, pg2 = [], [], []  # optimizer parameter groups
  116.     for k, v in model.named_modules():
  117.         if hasattr(v, 'bias') and isinstance(v.bias, nn.Parameter):
  118.             pg2.append(v.bias)  # biases
  119.         if isinstance(v, nn.BatchNorm2d):
  120.             pg0.append(v.weight)  # no decay
  121.         elif hasattr(v, 'weight') and isinstance(v.weight, nn.Parameter):
  122.             pg1.append(v.weight)  # apply decay
  123.  
  124.     if opt.adam:
  125.         optimizer = optim.Adam(pg0, lr=hyp['lr0'], betas=(hyp['momentum'], 0.999))  # adjust beta1 to momentum
  126.     else:
  127.         optimizer = optim.SGD(pg0, lr=hyp['lr0'], momentum=hyp['momentum'], nesterov=True)
  128.  
  129.     optimizer.add_param_group({'params': pg1, 'weight_decay': hyp['weight_decay']})  # add pg1 with weight_decay
  130.     optimizer.add_param_group({'params': pg2})  # add pg2 (biases)
  131.     logger.info('Optimizer groups: %g .bias, %g conv.weight, %g other' % (len(pg2), len(pg1), len(pg0)))
  132.     del pg0, pg1, pg2
  133.  
  134.     # Scheduler https://arxiv.org/pdf/1812.01187.pdf
  135.     # https://pytorch.org/docs/stable/_modules/torch/optim/lr_scheduler.html#OneCycleLR
  136.     if opt.linear_lr:
  137.         lf = lambda x: (1 - x / (epochs - 1)) * (1.0 - hyp['lrf']) + hyp['lrf']  # linear
  138.     else:
  139.         lf = one_cycle(1, hyp['lrf'], epochs)  # cosine 1->hyp['lrf']
  140.     scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf)
  141.     # plot_lr_scheduler(optimizer, scheduler, epochs)
  142.  
  143.     # EMA
  144.     ema = ModelEMA(model) if rank in [-1, 0] else None
  145.  
  146.     # Resume
  147.     start_epoch, best_fitness = 0, 0.0
  148.     if pretrained:
  149.         # Optimizer
  150.         if ckpt['optimizer'] is not None:
  151.             optimizer.load_state_dict(ckpt['optimizer'])
  152.             best_fitness = ckpt['best_fitness']
  153.  
  154.         # EMA
  155.         if ema and ckpt.get('ema'):
  156.             ema.ema.load_state_dict(ckpt['ema'].float().state_dict())
  157.             ema.updates = ckpt['updates']
  158.  
  159.         # Results
  160.         if ckpt.get('training_results') is not None:
  161.             results_file.write_text(ckpt['training_results'])  # write results.txt
  162.  
  163.         # Epochs
  164.         start_epoch = ckpt['epoch'] + 1
  165.         if opt.resume:
  166.             assert start_epoch > 0, '%s training to %g epochs is finished, nothing to resume.' % (weights, epochs)
  167.         if epochs < start_epoch:
  168.             logger.info('%s has been trained for %g epochs. Fine-tuning for %g additional epochs.' %
  169.                         (weights, ckpt['epoch'], epochs))
  170.             epochs += ckpt['epoch']  # finetune additional epochs
  171.  
  172.         del ckpt, state_dict
  173.  
  174.     # Image sizes
  175.     gs = max(int(model.stride.max()), 32)  # grid size (max stride)
  176.     nl = model.model[-1].nl  # number of detection layers (used for scaling hyp['obj'])
  177.     imgsz, imgsz_test = [check_img_size(x, gs) for x in opt.img_size]  # verify imgsz are gs-multiples
  178.  
  179.     # DP mode
  180.     if cuda and rank == -1 and torch.cuda.device_count() > 1:
  181.         model = torch.nn.DataParallel(model)
  182.  
  183.     # SyncBatchNorm
  184.     if opt.sync_bn and cuda and rank != -1:
  185.         model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model).to(device)
  186.         logger.info('Using SyncBatchNorm()')
  187.  
  188.     # Trainloader
  189.     dataloader, dataset = create_dataloader(train_path, imgsz, batch_size, gs, opt,
  190.                                             hyp=hyp, augment=True, cache=opt.cache_images, rect=opt.rect, rank=rank,
  191.                                             world_size=opt.world_size, workers=opt.workers,
  192.                                             image_weights=opt.image_weights, quad=opt.quad, prefix=colorstr('train: '))
  193.     mlc = np.concatenate(dataset.labels, 0)[:, 0].max()  # max label class
  194.     nb = len(dataloader)  # number of batches
  195.     assert mlc < nc, 'Label class %g exceeds nc=%g in %s. Possible class labels are 0-%g' % (mlc, nc, opt.data, nc - 1)
  196.  
  197.     # Process 0
  198.     if rank in [-1, 0]:
  199.         testloader = create_dataloader(test_path, imgsz_test, batch_size * 2, gs, opt,  # testloader
  200.                                        hyp=hyp, cache=opt.cache_images and not opt.notest, rect=True, rank=-1,
  201.                                        world_size=opt.world_size, workers=opt.workers,
  202.                                        pad=0.5, prefix=colorstr('val: '))[0]
  203.  
  204.         if not opt.resume:
  205.             labels = np.concatenate(dataset.labels, 0)
  206.             c = torch.tensor(labels[:, 0])  # classes
  207.             # cf = torch.bincount(c.long(), minlength=nc) + 1.  # frequency
  208.             # model._initialize_biases(cf.to(device))
  209.             if plots:
  210.                 plot_labels(labels, names, save_dir, loggers)
  211.                 if tb_writer:
  212.                     tb_writer.add_histogram('classes', c, 0)
  213.  
  214.             # Anchors
  215.             if not opt.noautoanchor:
  216.                 check_anchors(dataset, model=model, thr=hyp['anchor_t'], imgsz=imgsz)
  217.             model.half().float()  # pre-reduce anchor precision
  218.  
  219.     # DDP mode
  220.     if cuda and rank != -1:
  221.         model = DDP(model, device_ids=[opt.local_rank], output_device=opt.local_rank,
  222.                     # nn.MultiheadAttention incompatibility with DDP https://github.com/pytorch/pytorch/issues/26698
  223.                     find_unused_parameters=any(isinstance(layer, nn.MultiheadAttention) for layer in model.modules()))
  224.  
  225.     # Model parameters
  226.     hyp['box'] *= 3. / nl  # scale to layers
  227.     hyp['cls'] *= nc / 80. * 3. / nl  # scale to classes and layers
  228.     hyp['obj'] *= (imgsz / 640) ** 2 * 3. / nl  # scale to image size and layers
  229.     hyp['label_smoothing'] = opt.label_smoothing
  230.     model.nc = nc  # attach number of classes to model
  231.     model.hyp = hyp  # attach hyperparameters to model
  232.     model.gr = 1.0  # iou loss ratio (obj_loss = 1.0 or iou)
  233.     model.class_weights = labels_to_class_weights(dataset.labels, nc).to(device) * nc  # attach class weights
  234.     model.names = names
  235.  
  236.     # Start training
  237.     t0 = time.time()
  238.     nw = max(round(hyp['warmup_epochs'] * nb), 1000)  # number of warmup iterations, max(3 epochs, 1k iterations)
  239.     # nw = min(nw, (epochs - start_epoch) / 2 * nb)  # limit warmup to < 1/2 of training
  240.     maps = np.zeros(nc)  # mAP per class
  241.     results = (0, 0, 0, 0, 0, 0, 0)  # P, R, mAP@.5, mAP@.5-.95, val_loss(box, obj, cls)
  242.     scheduler.last_epoch = start_epoch - 1  # do not move
  243.     scaler = amp.GradScaler(enabled=cuda)
  244.     compute_loss = ComputeLoss(model)  # init loss class
  245.     logger.info(f'Image sizes {imgsz} train, {imgsz_test} test\n'
  246.                 f'Using {dataloader.num_workers} dataloader workers\n'
  247.                 f'Logging results to {save_dir}\n'
  248.                 f'Starting training for {epochs} epochs...')
  249.     for epoch in range(start_epoch, epochs):  # epoch ------------------------------------------------------------------
  250.         model.train()
  251.  
  252.         # Update image weights (optional)
  253.         if opt.image_weights:
  254.             # Generate indices
  255.             if rank in [-1, 0]:
  256.                 cw = model.class_weights.cpu().numpy() * (1 - maps) ** 2 / nc  # class weights
  257.                 iw = labels_to_image_weights(dataset.labels, nc=nc, class_weights=cw)  # image weights
  258.                 dataset.indices = random.choices(range(dataset.n), weights=iw, k=dataset.n)  # rand weighted idx
  259.             # Broadcast if DDP
  260.             if rank != -1:
  261.                 indices = (torch.tensor(dataset.indices) if rank == 0 else torch.zeros(dataset.n)).int()
  262.                 dist.broadcast(indices, 0)
  263.                 if rank != 0:
  264.                     dataset.indices = indices.cpu().numpy()
  265.  
  266.         # Update mosaic border
  267.         # b = int(random.uniform(0.25 * imgsz, 0.75 * imgsz + gs) // gs * gs)
  268.         # dataset.mosaic_border = [b - imgsz, -b]  # height, width borders
  269.  
  270.         mloss = torch.zeros(4, device=device)  # mean losses
  271.         if rank != -1:
  272.             dataloader.sampler.set_epoch(epoch)
  273.         pbar = enumerate(dataloader)
  274.         logger.info(('\n' + '%10s' * 8) % ('Epoch', 'gpu_mem', 'box', 'obj', 'cls', 'total', 'labels', 'img_size'))
  275.         if rank in [-1, 0]:
  276.             pbar = tqdm(pbar, total=nb)  # progress bar
  277.         optimizer.zero_grad()
  278.         for i, (imgs, targets, paths, _) in pbar:  # batch -------------------------------------------------------------
  279.             ni = i + nb * epoch  # number integrated batches (since train start)
  280.             imgs = imgs.to(device, non_blocking=True).float() / 255.0  # uint8 to float32, 0-255 to 0.0-1.0
  281.  
  282.             # Warmup
  283.             if ni <= nw:
  284.                 xi = [0, nw]  # x interp
  285.                 # model.gr = np.interp(ni, xi, [0.0, 1.0])  # iou loss ratio (obj_loss = 1.0 or iou)
  286.                 accumulate = max(1, np.interp(ni, xi, [1, nbs / total_batch_size]).round())
  287.                 for j, x in enumerate(optimizer.param_groups):
  288.                     # bias lr falls from 0.1 to lr0, all other lrs rise from 0.0 to lr0
  289.                     x['lr'] = np.interp(ni, xi, [hyp['warmup_bias_lr'] if j == 2 else 0.0, x['initial_lr'] * lf(epoch)])
  290.                     if 'momentum' in x:
  291.                         x['momentum'] = np.interp(ni, xi, [hyp['warmup_momentum'], hyp['momentum']])
  292.  
  293.             # Multi-scale
  294.             if opt.multi_scale:
  295.                 sz = random.randrange(imgsz * 0.5, imgsz * 1.5 + gs) // gs * gs  # size
  296.                 sf = sz / max(imgs.shape[2:])  # scale factor
  297.                 if sf != 1:
  298.                     ns = [math.ceil(x * sf / gs) * gs for x in imgs.shape[2:]]  # new shape (stretched to gs-multiple)
  299.                     imgs = F.interpolate(imgs, size=ns, mode='bilinear', align_corners=False)
  300.  
  301.             # Forward
  302.             with amp.autocast(enabled=cuda):
  303.                 pred = model(imgs)  # forward
  304.                 loss, loss_items = compute_loss(pred, targets.to(device))  # loss scaled by batch_size
  305.                 if rank != -1:
  306.                     loss *= opt.world_size  # gradient averaged between devices in DDP mode
  307.                 if opt.quad:
  308.                     loss *= 4.
  309.  
  310.             # Backward
  311.             scaler.scale(loss).backward()
  312.  
  313.             # Optimize
  314.             if ni % accumulate == 0:
  315.                 scaler.step(optimizer)  # optimizer.step
  316.                 scaler.update()
  317.                 optimizer.zero_grad()
  318.                 if ema:
  319.                     ema.update(model)
  320.  
  321.             # Print
  322.             if rank in [-1, 0]:
  323.                 mloss = (mloss * i + loss_items) / (i + 1)  # update mean losses
  324.                 mem = '%.3gG' % (torch.cuda.memory_reserved() / 1E9 if torch.cuda.is_available() else 0)  # (GB)
  325.                 s = ('%10s' * 2 + '%10.4g' * 6) % (
  326.                     '%g/%g' % (epoch, epochs - 1), mem, *mloss, targets.shape[0], imgs.shape[-1])
  327.                 pbar.set_description(s)
  328.  
  329.                 # Plot
  330.                 if plots and ni < 3:
  331.                     f = save_dir / f'train_batch{ni}.jpg'  # filename
  332.                     Thread(target=plot_images, args=(imgs, targets, paths, f), daemon=True).start()
  333.                     # if tb_writer:
  334.                     #     tb_writer.add_image(f, result, dataformats='HWC', global_step=epoch)
  335.                     #     tb_writer.add_graph(torch.jit.trace(model, imgs, strict=False), [])  # add model graph
  336.                 elif plots and ni == 10 and wandb_logger.wandb:
  337.                     wandb_logger.log({"Mosaics": [wandb_logger.wandb.Image(str(x), caption=x.name) for x in
  338.                                                   save_dir.glob('train*.jpg') if x.exists()]})
  339.  
  340.             # end batch ------------------------------------------------------------------------------------------------
  341.         # end epoch ----------------------------------------------------------------------------------------------------
  342.  
  343.         # Scheduler
  344.         lr = [x['lr'] for x in optimizer.param_groups]  # for tensorboard
  345.         scheduler.step()
  346.  
  347.         # DDP process 0 or single-GPU
  348.         if rank in [-1, 0]:
  349.             # mAP
  350.             ema.update_attr(model, include=['yaml', 'nc', 'hyp', 'gr', 'names', 'stride', 'class_weights'])
  351.             final_epoch = epoch + 1 == epochs
  352.             if not opt.notest or final_epoch:  # Calculate mAP
  353.                 wandb_logger.current_epoch = epoch + 1
  354.                 results, maps, times = test.test(data_dict,
  355.                                                  batch_size=batch_size * 2,
  356.                                                  imgsz=imgsz_test,
  357.                                                  model=ema.ema,
  358.                                                  single_cls=opt.single_cls,
  359.                                                  dataloader=testloader,
  360.                                                  save_dir=save_dir,
  361.                                                  verbose=nc < 50 and final_epoch,
  362.                                                  plots=plots and final_epoch,
  363.                                                  wandb_logger=wandb_logger,
  364.                                                  compute_loss=compute_loss,
  365.                                                  is_coco=is_coco)
  366.  
  367.             # Write
  368.             with open(results_file, 'a') as f:
  369.                 f.write(s + '%10.4g' * 7 % results + '\n')  # append metrics, val_loss
  370.             if len(opt.name) and opt.bucket:
  371.                 os.system('gsutil cp %s gs://%s/results/results%s.txt' % (results_file, opt.bucket, opt.name))
  372.  
  373.             # Log
  374.             tags = ['train/box_loss', 'train/obj_loss', 'train/cls_loss',  # train loss
  375.                     'metrics/precision', 'metrics/recall', 'metrics/mAP_0.5', 'metrics/mAP_0.5:0.95',
  376.                     'val/box_loss', 'val/obj_loss', 'val/cls_loss',  # val loss
  377.                     'x/lr0', 'x/lr1', 'x/lr2']  # params
  378.             for x, tag in zip(list(mloss[:-1]) + list(results) + lr, tags):
  379.                 if tb_writer:
  380.                     tb_writer.add_scalar(tag, x, epoch)  # tensorboard
  381.                 if wandb_logger.wandb:
  382.                     wandb_logger.log({tag: x})  # W&B
  383.  
  384.             # Update best mAP
  385.             fi = fitness(np.array(results).reshape(1, -1))  # weighted combination of [P, R, mAP@.5, mAP@.5-.95]
  386.             if fi > best_fitness:
  387.                 best_fitness = fi
  388.             wandb_logger.end_epoch(best_result=best_fitness == fi)
  389.  
  390.             # Save model
  391.             if (not opt.nosave) or (final_epoch and not opt.evolve):  # if save
  392.                 ckpt = {'epoch': epoch,
  393.                         'best_fitness': best_fitness,
  394.                         'training_results': results_file.read_text(),
  395.                         'model': deepcopy(model.module if is_parallel(model) else model).half(),
  396.                         'ema': deepcopy(ema.ema).half(),
  397.                         'updates': ema.updates,
  398.                         'optimizer': optimizer.state_dict(),
  399.                         'wandb_id': wandb_logger.wandb_run.id if wandb_logger.wandb else None}
  400.  
  401.                 # Save last, best and delete
  402.                 torch.save(ckpt, last)
  403.                 if best_fitness == fi:
  404.                     torch.save(ckpt, best)
  405.                 if wandb_logger.wandb:
  406.                     if ((epoch + 1) % opt.save_period == 0 and not final_epoch) and opt.save_period != -1:
  407.                         wandb_logger.log_model(
  408.                             last.parent, opt, epoch, fi, best_model=best_fitness == fi)
  409.                 del ckpt
  410.  
  411.         # end epoch ----------------------------------------------------------------------------------------------------
  412.     # end training
  413.     if rank in [-1, 0]:
  414.         # Plots
  415.         if plots:
  416.             plot_results(save_dir=save_dir)  # save as results.png
  417.             if wandb_logger.wandb:
  418.                 files = ['results.png', 'confusion_matrix.png', *[f'{x}_curve.png' for x in ('F1', 'PR', 'P', 'R')]]
  419.                 wandb_logger.log({"Results": [wandb_logger.wandb.Image(str(save_dir / f), caption=f) for f in files
  420.                                               if (save_dir / f).exists()]})
  421.         # Test best.pt
  422.         logger.info('%g epochs completed in %.3f hours.\n' % (epoch - start_epoch + 1, (time.time() - t0) / 3600))
  423.         if opt.data.endswith('coco.yaml') and nc == 80:  # if COCO
  424.             for m in (last, best) if best.exists() else (last):  # speed, mAP tests
  425.                 results, _, _ = test.test(opt.data,
  426.                                           batch_size=batch_size * 2,
  427.                                           imgsz=imgsz_test,
  428.                                           conf_thres=0.001,
  429.                                           iou_thres=0.7,
  430.                                           model=attempt_load(m, device).half(),
  431.                                           single_cls=opt.single_cls,
  432.                                           dataloader=testloader,
  433.                                           save_dir=save_dir,
  434.                                           save_json=True,
  435.                                           plots=False,
  436.                                           is_coco=is_coco)
  437.  
  438.         # Strip optimizers
  439.         final = best if best.exists() else last  # final model
  440.         for f in last, best:
  441.             if f.exists():
  442.                 strip_optimizer(f)  # strip optimizers
  443.         if opt.bucket:
  444.             os.system(f'gsutil cp {final} gs://{opt.bucket}/weights')  # upload
  445.         if wandb_logger.wandb and not opt.evolve:  # Log the stripped model
  446.             wandb_logger.wandb.log_artifact(str(final), type='model',
  447.                                             name='run_' + wandb_logger.wandb_run.id + '_model',
  448.                                             aliases=['last', 'best', 'stripped'])
  449.         wandb_logger.finish_run()
  450.     else:
  451.         dist.destroy_process_group()
  452.     torch.cuda.empty_cache()
  453.     return results
  454.  
  455.  
  456. if __name__ == '__main__':
  457.     parser = argparse.ArgumentParser()
  458.     parser.add_argument('--weights', type=str, default='weights/yolov5s.pt', help='initial weights path')
  459.     parser.add_argument('--cfg', type=str, default='models/yolov5s.yaml', help='model.yaml path')
  460.     parser.add_argument('--data', type=str, default='data/voc.yaml', help='data.yaml path')
  461.     parser.add_argument('--hyp', type=str, default='data/hyp.scratch.yaml', help='hyperparameters path')
  462.     parser.add_argument('--epochs', type=int, default=300)
  463.     parser.add_argument('--batch-size', type=int, default=16, help='total batch size for all GPUs')
  464.     parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='[train, test] image sizes')
  465.     parser.add_argument('--rect', action='store_true', help='rectangular training')
  466.     parser.add_argument('--resume', nargs='?', const=True, default=False, help='resume most recent training')
  467.     parser.add_argument('--nosave', action='store_true', help='only save final checkpoint')
  468.     parser.add_argument('--notest', action='store_true', help='only test final epoch')
  469.     parser.add_argument('--noautoanchor', action='store_true', help='disable autoanchor check')
  470.     parser.add_argument('--evolve', action='store_true', help='evolve hyperparameters')
  471.     parser.add_argument('--bucket', type=str, default='', help='gsutil bucket')
  472.     parser.add_argument('--cache-images', action='store_true', help='cache images for faster training')
  473.     parser.add_argument('--image-weights', action='store_true', help='use weighted image selection for training')
  474.     parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
  475.     parser.add_argument('--multi-scale', action='store_true', help='vary img-size +/- 50%%')
  476.     parser.add_argument('--single-cls', action='store_true', help='train multi-class data as single-class')
  477.     parser.add_argument('--adam', action='store_true', help='use torch.optim.Adam() optimizer')
  478.     parser.add_argument('--sync-bn', action='store_true', help='use SyncBatchNorm, only available in DDP mode')
  479.     parser.add_argument('--local_rank', type=int, default=-1, help='DDP parameter, do not modify')
  480.     parser.add_argument('--workers', type=int, default=0, help='maximum number of dataloader workers')
  481.     parser.add_argument('--project', default='runs/train', help='save to project/name')
  482.     parser.add_argument('--entity', default=None, help='W&B entity')
  483.     parser.add_argument('--name', default='exp', help='save to project/name')
  484.     parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
  485.     parser.add_argument('--quad', action='store_true', help='quad dataloader')
  486.     parser.add_argument('--linear-lr', action='store_true', help='linear LR')
  487.     parser.add_argument('--label-smoothing', type=float, default=0.0, help='Label smoothing epsilon')
  488.     parser.add_argument('--upload_dataset', action='store_true', help='Upload dataset as W&B artifact table')
  489.     parser.add_argument('--bbox_interval', type=int, default=-1, help='Set bounding-box image logging interval for W&B')
  490.     parser.add_argument('--save_period', type=int, default=-1, help='Log model after every "save_period" epoch')
  491.     parser.add_argument('--artifact_alias', type=str, default="latest", help='version of dataset artifact to be used')
  492.     opt = parser.parse_args()
  493.  
  494.     # Set DDP variables
  495.     opt.world_size = int(os.environ['WORLD_SIZE']) if 'WORLD_SIZE' in os.environ else 1
  496.     opt.global_rank = int(os.environ['RANK']) if 'RANK' in os.environ else -1
  497.     set_logging(opt.global_rank)
  498.     if opt.global_rank in [-1, 0]:
  499.         check_git_status()
  500.         check_requirements()
  501.  
  502.     # Resume
  503.     wandb_run = check_wandb_resume(opt)
  504.     if opt.resume and not wandb_run:  # resume an interrupted run
  505.         ckpt = opt.resume if isinstance(opt.resume, str) else get_latest_run()  # specified or most recent path
  506.         assert os.path.isfile(ckpt), 'ERROR: --resume checkpoint does not exist'
  507.         apriori = opt.global_rank, opt.local_rank
  508.         with open(Path(ckpt).parent.parent / 'opt.yaml') as f:
  509.             opt = argparse.Namespace(**yaml.load(f, Loader=yaml.SafeLoader))  # replace
  510.         opt.cfg, opt.weights, opt.resume, opt.batch_size, opt.global_rank, opt.local_rank = '', ckpt, True, opt.total_batch_size, *apriori  # reinstate
  511.         logger.info('Resuming training from %s' % ckpt)
  512.     else:
  513.         # opt.hyp = opt.hyp or ('hyp.finetune.yaml' if opt.weights else 'hyp.scratch.yaml')
  514.         opt.data, opt.cfg, opt.hyp = check_file(opt.data), check_file(opt.cfg), check_file(opt.hyp)  # check files
  515.         assert len(opt.cfg) or len(opt.weights), 'either --cfg or --weights must be specified'
  516.         opt.img_size.extend([opt.img_size[-1]] * (2 - len(opt.img_size)))  # extend to 2 sizes (train, test)
  517.         opt.name = 'evolve' if opt.evolve else opt.name
  518.         opt.save_dir = increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok | opt.evolve)  # increment run
  519.  
  520.     # DDP mode
  521.     opt.total_batch_size = opt.batch_size
  522.     device = select_device(opt.device, batch_size=opt.batch_size)
  523.     if opt.local_rank != -1:
  524.         assert torch.cuda.device_count() > opt.local_rank
  525.         torch.cuda.set_device(opt.local_rank)
  526.         device = torch.device('cuda', opt.local_rank)
  527.         dist.init_process_group(backend='nccl', init_method='env://')  # distributed backend
  528.         assert opt.batch_size % opt.world_size == 0, '--batch-size must be multiple of CUDA device count'
  529.         opt.batch_size = opt.total_batch_size // opt.world_size
  530.  
  531.     # Hyperparameters
  532.     with open(opt.hyp) as f:
  533.         hyp = yaml.load(f, Loader=yaml.SafeLoader)  # load hyps
  534.  
  535.     # Train
  536.     logger.info(opt)
  537.     if not opt.evolve:
  538.         tb_writer = None  # init loggers
  539.         if opt.global_rank in [-1, 0]:
  540.             prefix = colorstr('tensorboard: ')
  541.             logger.info(f"{prefix}Start with 'tensorboard --logdir {opt.project}', view at http://localhost:6006/")
  542.             tb_writer = SummaryWriter(opt.save_dir)  # Tensorboard
  543.         train(hyp, opt, device, tb_writer)
  544.  
  545.     # Evolve hyperparameters (optional)
  546.     else:
  547.         # Hyperparameter evolution metadata (mutation scale 0-1, lower_limit, upper_limit)
  548.         meta = {'lr0': (1, 1e-5, 1e-1),  # initial learning rate (SGD=1E-2, Adam=1E-3)
  549.                 'lrf': (1, 0.01, 1.0),  # final OneCycleLR learning rate (lr0 * lrf)
  550.                 'momentum': (0.3, 0.6, 0.98),  # SGD momentum/Adam beta1
  551.                 'weight_decay': (1, 0.0, 0.001),  # optimizer weight decay
  552.                 'warmup_epochs': (1, 0.0, 5.0),  # warmup epochs (fractions ok)
  553.                 'warmup_momentum': (1, 0.0, 0.95),  # warmup initial momentum
  554.                 'warmup_bias_lr': (1, 0.0, 0.2),  # warmup initial bias lr
  555.                 'box': (1, 0.02, 0.2),  # box loss gain
  556.                 'cls': (1, 0.2, 4.0),  # cls loss gain
  557.                 'cls_pw': (1, 0.5, 2.0),  # cls BCELoss positive_weight
  558.                 'obj': (1, 0.2, 4.0),  # obj loss gain (scale with pixels)
  559.                 'obj_pw': (1, 0.5, 2.0),  # obj BCELoss positive_weight
  560.                 'iou_t': (0, 0.1, 0.7),  # IoU training threshold
  561.                 'anchor_t': (1, 2.0, 8.0),  # anchor-multiple threshold
  562.                 'anchors': (2, 2.0, 10.0),  # anchors per output grid (0 to ignore)
  563.                 'fl_gamma': (0, 0.0, 2.0),  # focal loss gamma (efficientDet default gamma=1.5)
  564.                 'hsv_h': (1, 0.0, 0.1),  # image HSV-Hue augmentation (fraction)
  565.                 'hsv_s': (1, 0.0, 0.9),  # image HSV-Saturation augmentation (fraction)
  566.                 'hsv_v': (1, 0.0, 0.9),  # image HSV-Value augmentation (fraction)
  567.                 'degrees': (1, 0.0, 45.0),  # image rotation (+/- deg)
  568.                 'translate': (1, 0.0, 0.9),  # image translation (+/- fraction)
  569.                 'scale': (1, 0.0, 0.9),  # image scale (+/- gain)
  570.                 'shear': (1, 0.0, 10.0),  # image shear (+/- deg)
  571.                 'perspective': (0, 0.0, 0.001),  # image perspective (+/- fraction), range 0-0.001
  572.                 'flipud': (1, 0.0, 1.0),  # image flip up-down (probability)
  573.                 'fliplr': (0, 0.0, 1.0),  # image flip left-right (probability)
  574.                 'mosaic': (1, 0.0, 1.0),  # image mixup (probability)
  575.                 'mixup': (1, 0.0, 1.0)}  # image mixup (probability)
  576.  
  577.         assert opt.local_rank == -1, 'DDP mode not implemented for --evolve'
  578.         opt.notest, opt.nosave = True, True  # only test/save final epoch
  579.         # ei = [isinstance(x, (int, float)) for x in hyp.values()]  # evolvable indices
  580.         yaml_file = Path(opt.save_dir) / 'hyp_evolved.yaml'  # save best result here
  581.         if opt.bucket:
  582.             os.system('gsutil cp gs://%s/evolve.txt .' % opt.bucket)  # download evolve.txt if exists
  583.  
  584.         for _ in range(300):  # generations to evolve
  585.             if Path('evolve.txt').exists():  # if evolve.txt exists: select best hyps and mutate
  586.                 # Select parent(s)
  587.                 parent = 'single'  # parent selection method: 'single' or 'weighted'
  588.                 x = np.loadtxt('evolve.txt', ndmin=2)
  589.                 n = min(5, len(x))  # number of previous results to consider
  590.                 x = x[np.argsort(-fitness(x))][:n]  # top n mutations
  591.                 w = fitness(x) - fitness(x).min()  # weights
  592.                 if parent == 'single' or len(x) == 1:
  593.                     # x = x[random.randint(0, n - 1)]  # random selection
  594.                     x = x[random.choices(range(n), weights=w)[0]]  # weighted selection
  595.                 elif parent == 'weighted':
  596.                     x = (x * w.reshape(n, 1)).sum(0) / w.sum()  # weighted combination
  597.  
  598.                 # Mutate
  599.                 mp, s = 0.8, 0.2  # mutation probability, sigma
  600.                 npr = np.random
  601.                 npr.seed(int(time.time()))
  602.                 g = np.array([x[0] for x in meta.values()])  # gains 0-1
  603.                 ng = len(meta)
  604.                 v = np.ones(ng)
  605.                 while all(v == 1):  # mutate until a change occurs (prevent duplicates)
  606.                     v = (g * (npr.random(ng) < mp) * npr.randn(ng) * npr.random() * s + 1).clip(0.3, 3.0)
  607.                 for i, k in enumerate(hyp.keys()):  # plt.hist(v.ravel(), 300)
  608.                     hyp[k] = float(x[i + 7] * v[i])  # mutate
  609.  
  610.             # Constrain to limits
  611.             for k, v in meta.items():
  612.                 hyp[k] = max(hyp[k], v[1])  # lower limit
  613.                 hyp[k] = min(hyp[k], v[2])  # upper limit
  614.                 hyp[k] = round(hyp[k], 5)  # significant digits
  615.  
  616.             # Train mutation
  617.             results = train(hyp.copy(), opt, device)
  618.  
  619.             # Write mutation results
  620.             print_mutation(hyp.copy(), results, yaml_file, opt.bucket)
  621.  
  622.         # Plot results
  623.         plot_evolution(yaml_file)
  624.         print(f'Hyperparameter evolution complete. Best results saved as: {yaml_file}\n'
  625.               f'Command to train a new model with these hyperparameters: $ python train.py --hyp {yaml_file}')

声明：本文内容由网友自发贡献，不代表【wpsshop博客】立场，版权归原作者所有，本站不承担相应法律责任。如您发现有侵权的内容，请联系我们。转载请注明出处：https://www.wpsshop.cn/w/我家小花儿/article/detail/942374
推荐阅读
相关标签

Copyright © 2003-2013 www.wpsshop.cn 版权所有，并保留所有权利。

  

闽ICP备14008679号