基于深度学习的水果检测与识别系统（Python界面版，YOLOv5实现）
基于YOLOv5的实时水果识别系统与分类系统演示与介绍（Python+ui界面+训练代码+性能优化+实时检测识别）_哔哩哔哩_bilibili基于YOLOv5的实时水果识别系统与分类系统演示与介绍（Python+ui界面+训练代码+性能优化+实时检测识别）runs文件夹中，存放训练和评估的结果图import sysimport cv2# 添加一个关于界面# 窗口主类# 基本配置不动，然后只动第三个
·
目录
效果展示（完整源码在B站视频简介内）
基于YOLOv5的实时水果识别系统与分类系统演示与介绍（Python+ui界面+训练代码+性能优化+实时检测识别）_哔哩哔哩_bilibili
基于YOLOv5的实时水果识别系统与分类系统演示与介绍（Python+ui界面+训练代码+性能优化+实时检测识别）
（完整源码在B站视频简介内）
环境安装
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算法对常见水果进行检测和识别，实现对图片、视频和实时视频中的水果进行准确识别。博文详细阐述了算法原理，同时提供Python实现代码、训练数据集，以及基于PyQt的UI界面。通过YOLOv5实现对图像中存在的多个水果目标进行识别分类，用户可以在界面中选择各种水果图片、视频进行检测识别。本文旨在为相关领域的研究人员和新入门的朋友提供一个参考。
近年来，随着全球经济的发展，水果消费市场规模不断扩大，水果种类也日益丰富。水果检测与识别技术在农业生产、仓储物流、超市零售等领域具有重要的应用价值。传统的水果检测与识别方法主要依赖于人工识别，这种方法在一定程度上受到人力成本、识别效率和准确性等方面的限制。因此，开发一种高效、准确的自动化水果检测与识别系统具有重要的研究意义和实际价值。（本文的参考文献请见文末）
计算机视觉作为人工智能的一个重要分支，在目标检测和识别方面取得了显著的研究进展。特别是深度学习技术的发展，极大地推动了计算机视觉在水果检测与识别领域的应用。许多研究人员已经尝试利用深度学习技术进行水果检测与识别，取得了一定的成果[1]。然而，当前的研究仍然存在一定的局限性，如算法复杂度高、实时性差等问题。
在计算机视觉领域，已有多种深度学习算法被应用于目标检测和识别任务。例如，R-CNN[2]、Fast R-CNN[3]、Faster R-CNN[4]、SSD[5]、RetinaNet[6]等。这些算法在一定程度上提高了目标检测的准确性和速度。然而，这些算法仍然存在一定的局限性，如计算复杂度高、实时性差等。为了解决这些问题，研究人员提出了YOLO（You Only Look Once）算法[7]，该算法采用端到端的训练方式，能够在保持较高检测准确性的同时实现实时检测。YOLO算法自2016年首次提出以来，经历了多次改进，如YOLOv2[8]、YOLOv3[9]、YOLOv4[10]等，每个版本都在前一个版本的基础上提高了检测准确性和速度。针对水果检测与识别任务，有研究人员尝试将深度学习技术应用于该领域。例如，使用Faster R-CNN进行苹果检测[11]，采用Mask R-CNN实现对葡萄的实例分割[12]，以及利用SSD进行柑橘检测[13]等。这些研究表明深度学习技术在水果检测与识别任务上具有很大的前景。
在本博文中，我们提出了一种基于深度学习的水果检测与识别系统，该系统采用YOLOv5算法对常见水果进行检测和识别，实现对图片、视频和实时视频中的水果进行准确识别。YOLOv5[14]作为YOLO系列算法的最新版本，在保持实时性的同时，进一步提高了检测准确性。与其他目标检测算法相比，YOLOv5具有较高的性能和较低的计算复杂度，因此适合应用于水果检测与识别任务。
本文的主要贡献包括：（1）介绍了一种基于YOLOv5的水果检测与识别系统；（2）详细描述了算法原理，提供了Python实现代码以及训练数据集；（3）展示了基于PyQt的UI界面设计，并分析了训练和评估结果等实验。本文旨在为相关领域的研究人员和新入门的朋友提供一个参考。
核心源码（完整源码在B站视频简介内）
# -*- coding: UTF-8 -*-
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 = 'runs/train/yolov5s/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
# Load model
model = attempt_load(weights, 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:1px}"
"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('CSDN：https://blog.csdn.net/m0_68036862?type=blog\n哔哩哔哩：https://space.bilibili.com/1587012148?spm_id_from=333.1007.0.0\nCSDN搜索人工智能_SYBH') # 修改欢迎词语
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("人工智能_SYBH")
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 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"))
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)
# 生成读取视频对象
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（完整源码在B站视频简介内）
# YOLOv5 🚀 by Ultralytics, GPL-3.0 license
"""
Train a YOLOv5 model on a custom dataset.
Models and datasets download automatically from the latest YOLOv5 release.
Usage - Single-GPU training:
$ python train.py --data coco128.yaml --weights yolov5s.pt --img 640 # from pretrained (recommended)
$ python train.py --data coco128.yaml --weights '' --cfg yolov5s.yaml --img 640 # from scratch
Usage - Multi-GPU DDP training:
$ python -m torch.distributed.run --nproc_per_node 4 --master_port 1 train.py --data coco128.yaml --weights yolov5s.pt --img 640 --device 0,1,2,3
Models: https://github.com/ultralytics/yolov5/tree/master/models
Datasets: https://github.com/ultralytics/yolov5/tree/master/data
Tutorial: https://github.com/ultralytics/yolov5/wiki/Train-Custom-Data
"""
import argparse
import math
import os
import random
import sys
import time
from copy import deepcopy
from datetime import datetime
from pathlib import Path
import numpy as np
import torch
import torch.distributed as dist
import torch.nn as nn
import yaml
from torch.optim import lr_scheduler
from tqdm import tqdm
FILE = Path(__file__).resolve()
ROOT = FILE.parents[0] # YOLOv5 root directory
if str(ROOT) not in sys.path:
sys.path.append(str(ROOT)) # add ROOT to PATH
ROOT = Path(os.path.relpath(ROOT, Path.cwd())) # relative
import val as validate # for end-of-epoch mAP
from models.experimental import attempt_load
from models.yolo import Model
from utils.autoanchor import check_anchors
from utils.autobatch import check_train_batch_size
from utils.callbacks import Callbacks
from utils.dataloaders import create_dataloader
from utils.downloads import attempt_download, is_url
from utils.general import (LOGGER, check_amp, check_dataset, check_file, check_git_status, check_img_size,
check_requirements, check_suffix, check_yaml, colorstr, get_latest_run, increment_path,
init_seeds, intersect_dicts, labels_to_class_weights, labels_to_image_weights, methods,
one_cycle, print_args, print_mutation, strip_optimizer, yaml_save)
from utils.loggers import Loggers
from utils.loggers.comet.comet_utils import check_comet_resume
from utils.loggers.wandb.wandb_utils import check_wandb_resume
from utils.loss import ComputeLoss
from utils.metrics import fitness
from utils.plots import plot_evolve
from utils.torch_utils import (EarlyStopping, ModelEMA, de_parallel, select_device, smart_DDP, smart_optimizer,
smart_resume, torch_distributed_zero_first)
LOCAL_RANK = int(os.getenv('LOCAL_RANK', -1)) # https://pytorch.org/docs/stable/elastic/run.html
RANK = int(os.getenv('RANK', -1))
WORLD_SIZE = int(os.getenv('WORLD_SIZE', 1))
def train(hyp, opt, device, callbacks): # hyp is path/to/hyp.yaml or hyp dictionary
save_dir, epochs, batch_size, weights, single_cls, evolve, data, cfg, resume, noval, nosave, workers, freeze = \
Path(opt.save_dir), opt.epochs, opt.batch_size, opt.weights, opt.single_cls, opt.evolve, opt.data, opt.cfg, \
opt.resume, opt.noval, opt.nosave, opt.workers, opt.freeze
callbacks.run('on_pretrain_routine_start')
# Directories
w = save_dir / 'weights' # weights dir
(w.parent if evolve else w).mkdir(parents=True, exist_ok=True) # make dir
last, best = w / 'last.pt', w / 'best.pt'
# Hyperparameters
if isinstance(hyp, str):
with open(hyp, errors='ignore') as f:
hyp = yaml.safe_load(f) # load hyps dict
LOGGER.info(colorstr('hyperparameters: ') + ', '.join(f'{k}={v}' for k, v in hyp.items()))
opt.hyp = hyp.copy() # for saving hyps to checkpoints
# Save run settings
if not evolve:
yaml_save(save_dir / 'hyp.yaml', hyp)
yaml_save(save_dir / 'opt.yaml', vars(opt))
# Loggers
data_dict = None
if RANK in {-1, 0}:
loggers = Loggers(save_dir, weights, opt, hyp, LOGGER) # loggers instance
# Register actions
for k in methods(loggers):
callbacks.register_action(k, callback=getattr(loggers, k))
# Process custom dataset artifact link
data_dict = loggers.remote_dataset
if resume: # If resuming runs from remote artifact
weights, epochs, hyp, batch_size = opt.weights, opt.epochs, opt.hyp, opt.batch_size
# Config
plots = not evolve and not opt.noplots # create plots
cuda = device.type != 'cpu'
init_seeds(opt.seed + 1 + RANK, deterministic=True)
with torch_distributed_zero_first(LOCAL_RANK):
data_dict = data_dict or check_dataset(data) # check if None
train_path, val_path = data_dict['train'], data_dict['val']
nc = 1 if single_cls else int(data_dict['nc']) # number of classes
names = {0: 'item'} if single_cls and len(data_dict['names']) != 1 else data_dict['names'] # class names
is_coco = isinstance(val_path, str) and val_path.endswith('coco/val2017.txt') # COCO dataset
# Model
check_suffix(weights, '.pt') # check weights
pretrained = weights.endswith('.pt')
if pretrained:
with torch_distributed_zero_first(LOCAL_RANK):
weights = attempt_download(weights) # download if not found locally
ckpt = torch.load(weights, map_location='cpu') # load checkpoint to CPU to avoid CUDA memory leak
model = Model(cfg or ckpt['model'].yaml, ch=3, nc=nc, anchors=hyp.get('anchors')).to(device) # create
exclude = ['anchor'] if (cfg or hyp.get('anchors')) and not resume else [] # exclude keys
csd = ckpt['model'].float().state_dict() # checkpoint state_dict as FP32
csd = intersect_dicts(csd, model.state_dict(), exclude=exclude) # intersect
model.load_state_dict(csd, strict=False) # load
LOGGER.info(f'Transferred {len(csd)}/{len(model.state_dict())} items from {weights}') # report
else:
model = Model(cfg, ch=3, nc=nc, anchors=hyp.get('anchors')).to(device) # create
amp = check_amp(model) # check AMP
# Freeze
freeze = [f'model.{x}.' for x in (freeze if len(freeze) > 1 else range(freeze[0]))] # layers to freeze
for k, v in model.named_parameters():
v.requires_grad = True # train all layers
# v.register_hook(lambda x: torch.nan_to_num(x)) # NaN to 0 (commented for erratic training results)
if any(x in k for x in freeze):
LOGGER.info(f'freezing {k}')
v.requires_grad = False
# Image size
gs = max(int(model.stride.max()), 32) # grid size (max stride)
imgsz = check_img_size(opt.imgsz, gs, floor=gs * 2) # verify imgsz is gs-multiple
# Batch size
if RANK == -1 and batch_size == -1: # single-GPU only, estimate best batch size
batch_size = check_train_batch_size(model, imgsz, amp)
loggers.on_params_update({"batch_size": batch_size})
# Optimizer
nbs = 64 # nominal batch size
accumulate = max(round(nbs / batch_size), 1) # accumulate loss before optimizing
hyp['weight_decay'] *= batch_size * accumulate / nbs # scale weight_decay
optimizer = smart_optimizer(model, opt.optimizer, hyp['lr0'], hyp['momentum'], hyp['weight_decay'])
# Scheduler
if opt.cos_lr:
lf = one_cycle(1, hyp['lrf'], epochs) # cosine 1->hyp['lrf']
else:
lf = lambda x: (1 - x / epochs) * (1.0 - hyp['lrf']) + hyp['lrf'] # linear
scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf) # plot_lr_scheduler(optimizer, scheduler, epochs)
# EMA
ema = ModelEMA(model) if RANK in {-1, 0} else None
# Resume
best_fitness, start_epoch = 0.0, 0
if pretrained:
if resume:
best_fitness, start_epoch, epochs = smart_resume(ckpt, optimizer, ema, weights, epochs, resume)
del ckpt, csd
# DP mode
if cuda and RANK == -1 and torch.cuda.device_count() > 1:
LOGGER.warning('WARNING: DP not recommended, use torch.distributed.run for best DDP Multi-GPU results.\n'
'See Multi-GPU Tutorial at https://github.com/ultralytics/yolov5/issues/475 to get started.')
model = torch.nn.DataParallel(model)
# SyncBatchNorm
if opt.sync_bn and cuda and RANK != -1:
model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model).to(device)
LOGGER.info('Using SyncBatchNorm()')
# Trainloader
train_loader, dataset = create_dataloader(train_path,
imgsz,
batch_size // WORLD_SIZE,
gs,
single_cls,
hyp=hyp,
augment=True,
cache=None if opt.cache == 'val' else opt.cache,
rect=opt.rect,
rank=LOCAL_RANK,
workers=workers,
image_weights=opt.image_weights,
quad=opt.quad,
prefix=colorstr('train: '),
shuffle=True)
labels = np.concatenate(dataset.labels, 0)
mlc = int(labels[:, 0].max()) # max label class
assert mlc < nc, f'Label class {mlc} exceeds nc={nc} in {data}. Possible class labels are 0-{nc - 1}'
# Process 0
if RANK in {-1, 0}:
val_loader = create_dataloader(val_path,
imgsz,
batch_size // WORLD_SIZE * 2,
gs,
single_cls,
hyp=hyp,
cache=None if noval else opt.cache,
rect=True,
rank=-1,
workers=workers * 2,
pad=0.5,
prefix=colorstr('val: '))[0]
if not resume:
if not opt.noautoanchor:
check_anchors(dataset, model=model, thr=hyp['anchor_t'], imgsz=imgsz) # run AutoAnchor
model.half().float() # pre-reduce anchor precision
callbacks.run('on_pretrain_routine_end', labels, names)
# DDP mode
if cuda and RANK != -1:
model = smart_DDP(model)
# Model attributes
nl = de_parallel(model).model[-1].nl # number of detection layers (to scale hyps)
hyp['box'] *= 3 / nl # scale to layers
hyp['cls'] *= nc / 80 * 3 / nl # scale to classes and layers
hyp['obj'] *= (imgsz / 640) ** 2 * 3 / nl # scale to image size and layers
hyp['label_smoothing'] = opt.label_smoothing
model.nc = nc # attach number of classes to model
model.hyp = hyp # attach hyperparameters to model
model.class_weights = labels_to_class_weights(dataset.labels, nc).to(device) * nc # attach class weights
model.names = names
# Start training
t0 = time.time()
nb = len(train_loader) # number of batches
nw = max(round(hyp['warmup_epochs'] * nb), 100) # number of warmup iterations, max(3 epochs, 100 iterations)
# nw = min(nw, (epochs - start_epoch) / 2 * nb) # limit warmup to < 1/2 of training
last_opt_step = -1
maps = np.zeros(nc) # mAP per class
results = (0, 0, 0, 0, 0, 0, 0) # P, R, mAP@.5, mAP@.5-.95, val_loss(box, obj, cls)
scheduler.last_epoch = start_epoch - 1 # do not move
scaler = torch.cuda.amp.GradScaler(enabled=amp)
stopper, stop = EarlyStopping(patience=opt.patience), False
compute_loss = ComputeLoss(model) # init loss class
callbacks.run('on_train_start')
LOGGER.info(f'Image sizes {imgsz} train, {imgsz} val\n'
f'Using {train_loader.num_workers * WORLD_SIZE} dataloader workers\n'
f"Logging results to {colorstr('bold', save_dir)}\n"
f'Starting training for {epochs} epochs...')
for epoch in range(start_epoch, epochs): # epoch ------------------------------------------------------------------
callbacks.run('on_train_epoch_start')
model.train()
# Update image weights (optional, single-GPU only)
if opt.image_weights:
cw = model.class_weights.cpu().numpy() * (1 - maps) ** 2 / nc # class weights
iw = labels_to_image_weights(dataset.labels, nc=nc, class_weights=cw) # image weights
dataset.indices = random.choices(range(dataset.n), weights=iw, k=dataset.n) # rand weighted idx
# Update mosaic border (optional)
# b = int(random.uniform(0.25 * imgsz, 0.75 * imgsz + gs) // gs * gs)
# dataset.mosaic_border = [b - imgsz, -b] # height, width borders
mloss = torch.zeros(3, device=device) # mean losses
if RANK != -1:
train_loader.sampler.set_epoch(epoch)
pbar = enumerate(train_loader)
LOGGER.info(('\n' + '%11s' * 7) % ('Epoch', 'GPU_mem', 'box_loss', 'obj_loss', 'cls_loss', 'Instances', 'Size'))
if RANK in {-1, 0}:
pbar = tqdm(pbar, total=nb, bar_format='{l_bar}{bar:10}{r_bar}{bar:-10b}') # progress bar
optimizer.zero_grad()
for i, (imgs, targets, paths, _) in pbar: # batch -------------------------------------------------------------
callbacks.run('on_train_batch_start')
ni = i + nb * epoch # number integrated batches (since train start)
imgs = imgs.to(device, non_blocking=True).float() / 255 # uint8 to float32, 0-255 to 0.0-1.0
# Warmup
if ni <= nw:
xi = [0, nw] # x interp
# compute_loss.gr = np.interp(ni, xi, [0.0, 1.0]) # iou loss ratio (obj_loss = 1.0 or iou)
accumulate = max(1, np.interp(ni, xi, [1, nbs / batch_size]).round())
for j, x in enumerate(optimizer.param_groups):
# bias lr falls from 0.1 to lr0, all other lrs rise from 0.0 to lr0
x['lr'] = np.interp(ni, xi, [hyp['warmup_bias_lr'] if j == 0 else 0.0, x['initial_lr'] * lf(epoch)])
if 'momentum' in x:
x['momentum'] = np.interp(ni, xi, [hyp['warmup_momentum'], hyp['momentum']])
# Multi-scale
if opt.multi_scale:
sz = random.randrange(imgsz * 0.5, imgsz * 1.5 + gs) // gs * gs # size
sf = sz / max(imgs.shape[2:]) # scale factor
if sf != 1:
ns = [math.ceil(x * sf / gs) * gs for x in imgs.shape[2:]] # new shape (stretched to gs-multiple)
imgs = nn.functional.interpolate(imgs, size=ns, mode='bilinear', align_corners=False)
# Forward
with torch.cuda.amp.autocast(amp):
pred = model(imgs) # forward
loss, loss_items = compute_loss(pred, targets.to(device)) # loss scaled by batch_size
if RANK != -1:
loss *= WORLD_SIZE # gradient averaged between devices in DDP mode
if opt.quad:
loss *= 4.
# Backward
scaler.scale(loss).backward()
# Optimize - https://pytorch.org/docs/master/notes/amp_examples.html
if ni - last_opt_step >= accumulate:
scaler.unscale_(optimizer) # unscale gradients
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=10.0) # clip gradients
scaler.step(optimizer) # optimizer.step
scaler.update()
optimizer.zero_grad()
if ema:
ema.update(model)
last_opt_step = ni
# Log
if RANK in {-1, 0}:
mloss = (mloss * i + loss_items) / (i + 1) # update mean losses
mem = f'{torch.cuda.memory_reserved() / 1E9 if torch.cuda.is_available() else 0:.3g}G' # (GB)
pbar.set_description(('%11s' * 2 + '%11.4g' * 5) %
(f'{epoch}/{epochs - 1}', mem, *mloss, targets.shape[0], imgs.shape[-1]))
callbacks.run('on_train_batch_end', model, ni, imgs, targets, paths, list(mloss))
if callbacks.stop_training:
return
# end batch ------------------------------------------------------------------------------------------------
# Scheduler
lr = [x['lr'] for x in optimizer.param_groups] # for loggers
scheduler.step()
if RANK in {-1, 0}:
# mAP
callbacks.run('on_train_epoch_end', epoch=epoch)
ema.update_attr(model, include=['yaml', 'nc', 'hyp', 'names', 'stride', 'class_weights'])
final_epoch = (epoch + 1 == epochs) or stopper.possible_stop
if not noval or final_epoch: # Calculate mAP
results, maps, _ = validate.run(data_dict,
batch_size=batch_size // WORLD_SIZE * 2,
imgsz=imgsz,
half=amp,
model=ema.ema,
single_cls=single_cls,
dataloader=val_loader,
save_dir=save_dir,
plots=False,
callbacks=callbacks,
compute_loss=compute_loss)
# Update best mAP
fi = fitness(np.array(results).reshape(1, -1)) # weighted combination of [P, R, mAP@.5, mAP@.5-.95]
stop = stopper(epoch=epoch, fitness=fi) # early stop check
if fi > best_fitness:
best_fitness = fi
log_vals = list(mloss) + list(results) + lr
callbacks.run('on_fit_epoch_end', log_vals, epoch, best_fitness, fi)
# Save model
if (not nosave) or (final_epoch and not evolve): # if save
ckpt = {
'epoch': epoch,
'best_fitness': best_fitness,
'model': deepcopy(de_parallel(model)).half(),
'ema': deepcopy(ema.ema).half(),
'updates': ema.updates,
'optimizer': optimizer.state_dict(),
'wandb_id': loggers.wandb.wandb_run.id if loggers.wandb else None,
'opt': vars(opt),
'date': datetime.now().isoformat()}
# Save last, best and delete
torch.save(ckpt, last)
if best_fitness == fi:
torch.save(ckpt, best)
if opt.save_period > 0 and epoch % opt.save_period == 0:
torch.save(ckpt, w / f'epoch{epoch}.pt')
del ckpt
callbacks.run('on_model_save', last, epoch, final_epoch, best_fitness, fi)
# EarlyStopping
if RANK != -1: # if DDP training
broadcast_list = [stop if RANK == 0 else None]
dist.broadcast_object_list(broadcast_list, 0) # broadcast 'stop' to all ranks
if RANK != 0:
stop = broadcast_list[0]
if stop:
break # must break all DDP ranks
# end epoch ----------------------------------------------------------------------------------------------------
# end training -----------------------------------------------------------------------------------------------------
if RANK in {-1, 0}:
LOGGER.info(f'\n{epoch - start_epoch + 1} epochs completed in {(time.time() - t0) / 3600:.3f} hours.')
for f in last, best:
if f.exists():
strip_optimizer(f) # strip optimizers
if f is best:
LOGGER.info(f'\nValidating {f}...')
results, _, _ = validate.run(
data_dict,
batch_size=batch_size // WORLD_SIZE * 2,
imgsz=imgsz,
model=attempt_load(f, device).half(),
iou_thres=0.65 if is_coco else 0.60, # best pycocotools at iou 0.65
single_cls=single_cls,
dataloader=val_loader,
save_dir=save_dir,
save_json=is_coco,
verbose=True,
plots=plots,
callbacks=callbacks,
compute_loss=compute_loss) # val best model with plots
if is_coco:
callbacks.run('on_fit_epoch_end', list(mloss) + list(results) + lr, epoch, best_fitness, fi)
callbacks.run('on_train_end', last, best, epoch, results)
torch.cuda.empty_cache()
return results
def parse_opt(known=False):
parser = argparse.ArgumentParser()
parser.add_argument('--weights', type=str, default=ROOT / 'weights/yolov5s.pt', help='initial weights path')
parser.add_argument('--cfg', type=str, default='models/ShuffleNetV2.yaml', help='model.yaml path')
parser.add_argument('--data', type=str, default=ROOT / 'data/VOC.yaml', help='dataset.yaml path')
parser.add_argument('--hyp', type=str, default=ROOT / 'data/hyps/hyp.scratch-low.yaml', help='hyperparameters path')
parser.add_argument('--epochs', type=int, default=100, help='total training epochs')
parser.add_argument('--batch-size', type=int, default=8, help='total batch size for all GPUs, -1 for autobatch')
parser.add_argument('--imgsz', '--img', '--img-size', type=int, default=640, help='train, val image size (pixels)')
parser.add_argument('--rect', action='store_true', help='rectangular training')
parser.add_argument('--resume', nargs='?', const=True, default=False, help='resume most recent training')
parser.add_argument('--nosave', action='store_true', help='only save final checkpoint')
parser.add_argument('--noval', action='store_true', help='only validate final epoch')
parser.add_argument('--noautoanchor', action='store_true', help='disable AutoAnchor')
parser.add_argument('--noplots', action='store_true', help='save no plot files')
parser.add_argument('--evolve', type=int, nargs='?', const=300, help='evolve hyperparameters for x generations')
parser.add_argument('--bucket', type=str, default='', help='gsutil bucket')
parser.add_argument('--cache', type=str, nargs='?', const='ram', help='--cache images in "ram" (default) or "disk"')
parser.add_argument('--image-weights', action='store_true', help='use weighted image selection for training')
parser.add_argument('--device', default='0', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
parser.add_argument('--multi-scale', action='store_true', help='vary img-size +/- 50%%')
parser.add_argument('--single-cls', action='store_true', help='train multi-class data as single-class')
parser.add_argument('--optimizer', type=str, choices=['SGD', 'Adam', 'AdamW'], default='SGD', help='optimizer')
parser.add_argument('--sync-bn', action='store_true', help='use SyncBatchNorm, only available in DDP mode')
parser.add_argument('--workers', type=int, default=2, help='max dataloader workers (per RANK in DDP mode)')
parser.add_argument('--project', default=ROOT / 'runs/train', help='save to project/name')
parser.add_argument('--name', default='exp', help='save to project/name')
parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
parser.add_argument('--quad', action='store_true', help='quad dataloader')
parser.add_argument('--cos-lr', action='store_true', help='cosine LR scheduler')
parser.add_argument('--label-smoothing', type=float, default=0.0, help='Label smoothing epsilon')
parser.add_argument('--patience', type=int, default=100, help='EarlyStopping patience (epochs without improvement)')
parser.add_argument('--freeze', nargs='+', type=int, default=[0], help='Freeze layers: backbone=10, first3=0 1 2')
parser.add_argument('--save-period', type=int, default=-1, help='Save checkpoint every x epochs (disabled if < 1)')
parser.add_argument('--seed', type=int, default=0, help='Global training seed')
parser.add_argument('--local_rank', type=int, default=-1, help='Automatic DDP Multi-GPU argument, do not modify')
# Logger arguments
parser.add_argument('--entity', default=None, help='Entity')
parser.add_argument('--upload_dataset', nargs='?', const=True, default=False, help='Upload data, "val" option')
parser.add_argument('--bbox_interval', type=int, default=-1, help='Set bounding-box image logging interval')
parser.add_argument('--artifact_alias', type=str, default='latest', help='Version of dataset artifact to use')
return parser.parse_known_args()[0] if known else parser.parse_args()
def main(opt, callbacks=Callbacks()):
# Checks
if RANK in {-1, 0}:
print_args(vars(opt))
check_git_status()
check_requirements()
# Resume (from specified or most recent last.pt)
if opt.resume and not check_wandb_resume(opt) and not check_comet_resume(opt) and not opt.evolve:
last = Path(check_file(opt.resume) if isinstance(opt.resume, str) else get_latest_run())
opt_yaml = last.parent.parent / 'opt.yaml' # train options yaml
opt_data = opt.data # original dataset
if opt_yaml.is_file():
with open(opt_yaml, errors='ignore') as f:
d = yaml.safe_load(f)
else:
d = torch.load(last, map_location='cpu')['opt']
opt = argparse.Namespace(**d) # replace
opt.cfg, opt.weights, opt.resume = '', str(last), True # reinstate
if is_url(opt_data):
opt.data = check_file(opt_data) # avoid HUB resume auth timeout
else:
opt.data, opt.cfg, opt.hyp, opt.weights, opt.project = \
check_file(opt.data), check_yaml(opt.cfg), check_yaml(opt.hyp), str(opt.weights), str(opt.project) # checks
assert len(opt.cfg) or len(opt.weights), 'either --cfg or --weights must be specified'
if opt.evolve:
if opt.project == str(ROOT / 'runs/train'): # if default project name, rename to runs/evolve
opt.project = str(ROOT / 'runs/evolve')
opt.exist_ok, opt.resume = opt.resume, False # pass resume to exist_ok and disable resume
if opt.name == 'cfg':
opt.name = Path(opt.cfg).stem # use model.yaml as name
opt.save_dir = str(increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok))
# DDP mode
device = select_device(opt.device, batch_size=opt.batch_size)
if LOCAL_RANK != -1:
msg = 'is not compatible with YOLOv5 Multi-GPU DDP training'
assert not opt.image_weights, f'--image-weights {msg}'
assert not opt.evolve, f'--evolve {msg}'
assert opt.batch_size != -1, f'AutoBatch with --batch-size -1 {msg}, please pass a valid --batch-size'
assert opt.batch_size % WORLD_SIZE == 0, f'--batch-size {opt.batch_size} must be multiple of WORLD_SIZE'
assert torch.cuda.device_count() > LOCAL_RANK, 'insufficient CUDA devices for DDP command'
torch.cuda.set_device(LOCAL_RANK)
device = torch.device('cuda', LOCAL_RANK)
dist.init_process_group(backend="nccl" if dist.is_nccl_available() else "gloo")
# Train
if not opt.evolve:
train(opt.hyp, opt, device, callbacks)
# Evolve hyperparameters (optional)
else:
# Hyperparameter evolution metadata (mutation scale 0-1, lower_limit, upper_limit)
meta = {
'lr0': (1, 1e-5, 1e-1), # initial learning rate (SGD=1E-2, Adam=1E-3)
'lrf': (1, 0.01, 1.0), # final OneCycleLR learning rate (lr0 * lrf)
'momentum': (0.3, 0.6, 0.98), # SGD momentum/Adam beta1
'weight_decay': (1, 0.0, 0.001), # optimizer weight decay
'warmup_epochs': (1, 0.0, 5.0), # warmup epochs (fractions ok)
'warmup_momentum': (1, 0.0, 0.95), # warmup initial momentum
'warmup_bias_lr': (1, 0.0, 0.2), # warmup initial bias lr
'box': (1, 0.02, 0.2), # box loss gain
'cls': (1, 0.2, 4.0), # cls loss gain
'cls_pw': (1, 0.5, 2.0), # cls BCELoss positive_weight
'obj': (1, 0.2, 4.0), # obj loss gain (scale with pixels)
'obj_pw': (1, 0.5, 2.0), # obj BCELoss positive_weight
'iou_t': (0, 0.1, 0.7), # IoU training threshold
'anchor_t': (1, 2.0, 8.0), # anchor-multiple threshold
'anchors': (2, 2.0, 10.0), # anchors per output grid (0 to ignore)
'fl_gamma': (0, 0.0, 2.0), # focal loss gamma (efficientDet default gamma=1.5)
'hsv_h': (1, 0.0, 0.1), # image HSV-Hue augmentation (fraction)
'hsv_s': (1, 0.0, 0.9), # image HSV-Saturation augmentation (fraction)
'hsv_v': (1, 0.0, 0.9), # image HSV-Value augmentation (fraction)
'degrees': (1, 0.0, 45.0), # image rotation (+/- deg)
'translate': (1, 0.0, 0.9), # image translation (+/- fraction)
'scale': (1, 0.0, 0.9), # image scale (+/- gain)
'shear': (1, 0.0, 10.0), # image shear (+/- deg)
'perspective': (0, 0.0, 0.001), # image perspective (+/- fraction), range 0-0.001
'flipud': (1, 0.0, 1.0), # image flip up-down (probability)
'fliplr': (0, 0.0, 1.0), # image flip left-right (probability)
'mosaic': (1, 0.0, 1.0), # image mixup (probability)
'mixup': (1, 0.0, 1.0), # image mixup (probability)
'copy_paste': (1, 0.0, 1.0)} # segment copy-paste (probability)
with open(opt.hyp, errors='ignore') as f:
hyp = yaml.safe_load(f) # load hyps dict
if 'anchors' not in hyp: # anchors commented in hyp.yaml
hyp['anchors'] = 3
if opt.noautoanchor:
del hyp['anchors'], meta['anchors']
opt.noval, opt.nosave, save_dir = True, True, Path(opt.save_dir) # only val/save final epoch
# ei = [isinstance(x, (int, float)) for x in hyp.values()] # evolvable indices
evolve_yaml, evolve_csv = save_dir / 'hyp_evolve.yaml', save_dir / 'evolve.csv'
if opt.bucket:
os.system(f'gsutil cp gs://{opt.bucket}/evolve.csv {evolve_csv}') # download evolve.csv if exists
for _ in range(opt.evolve): # generations to evolve
if evolve_csv.exists(): # if evolve.csv exists: select best hyps and mutate
# Select parent(s)
parent = 'single' # parent selection method: 'single' or 'weighted'
x = np.loadtxt(evolve_csv, ndmin=2, delimiter=',', skiprows=1)
n = min(5, len(x)) # number of previous results to consider
x = x[np.argsort(-fitness(x))][:n] # top n mutations
w = fitness(x) - fitness(x).min() + 1E-6 # weights (sum > 0)
if parent == 'single' or len(x) == 1:
# x = x[random.randint(0, n - 1)] # random selection
x = x[random.choices(range(n), weights=w)[0]] # weighted selection
elif parent == 'weighted':
x = (x * w.reshape(n, 1)).sum(0) / w.sum() # weighted combination
# Mutate
mp, s = 0.8, 0.2 # mutation probability, sigma
npr = np.random
npr.seed(int(time.time()))
g = np.array([meta[k][0] for k in hyp.keys()]) # gains 0-1
ng = len(meta)
v = np.ones(ng)
while all(v == 1): # mutate until a change occurs (prevent duplicates)
v = (g * (npr.random(ng) < mp) * npr.randn(ng) * npr.random() * s + 1).clip(0.3, 3.0)
for i, k in enumerate(hyp.keys()): # plt.hist(v.ravel(), 300)
hyp[k] = float(x[i + 7] * v[i]) # mutate
# Constrain to limits
for k, v in meta.items():
hyp[k] = max(hyp[k], v[1]) # lower limit
hyp[k] = min(hyp[k], v[2]) # upper limit
hyp[k] = round(hyp[k], 5) # significant digits
# Train mutation
results = train(hyp.copy(), opt, device, callbacks)
callbacks = Callbacks()
# Write mutation results
print_mutation(results, hyp.copy(), save_dir, opt.bucket)
# Plot results
plot_evolve(evolve_csv)
LOGGER.info(f'Hyperparameter evolution finished {opt.evolve} generations\n'
f"Results saved to {colorstr('bold', save_dir)}\n"
f'Usage example: $ python train.py --hyp {evolve_yaml}')
def run(**kwargs):
# Usage: import train; train.run(data='coco128.yaml', imgsz=320, weights='yolov5m.pt')
opt = parse_opt(True)
for k, v in kwargs.items():
setattr(opt, k, v)
main(opt)
return opt
if __name__ == "__main__":
opt = parse_opt()
main(opt)
技术共进，成长同行——讯飞AI开发者社区
更多推荐
33
0- 0
已为社区贡献24条内容
所有评论(0)