往期回顾:YOLOv5源码解读1.0-目录_汉卿HanQ的博客-CSDN博客
目录
7.整体代码
介绍完整体目录后,要开始学习代码啦(头疼!!!!)这里首先介绍detect.py,当我们导入好YOLO后,验证是否可以跑起来,第一件事就是运行detect.py,那么我们接下来看看这个代码逻辑吧!(因为csdn原因,推荐粘贴全部代码,然后按模块逐行学习!)
1.导入Python库
#----------------------------------1.导入Python库----------------------------------
"""它是一个用于命令项选项与参数解析的模块,通过在程序中定义好我们需要的参数,argparse将会从sys.argv中解析出这些参数,并自动生成帮助和使用信息"""
import argparse # 解析命令行参数的库
"""它提供了多种操作系统的接口。通过os模块提供的操作系统接口,我们可以对操作系统里文件、终端、进程等进行操作"""
import os # 与操作系统进行交互的文件库 包含文件路径操作与解析
"""它是与python解释器交互的一个接口,该模块提供对解释器使用或维护的一些变量的访问和获取,它提供了许多函数和变量来处理Python 运行时环境的不同部分"""
import sys # sys模块包含了与python解释器和它的环境有关的函数。
"""这个库提供了一种面向对象的方式来与文件系统交互,可以让代码更简洁、更易读"""
from pathlib import Path # Path能够更加方便得对字符串路径进行处理
"""opencv库"""
import cv2 # opencv库
"""这是主要的Pytorch库。它提供了构建、训练和评估神经网络的工具"""
import torch # pytorch 深度学习库
"""它提供了一个接口,用于使用cuDNN库,在NVIDIA GPU上高效地进行深度学习。cudnn模块是一个Pytorch库的扩展"""
import torch.backends.cudnn as cudnn # 让内置的cudnn的 auto-tuner 自动寻找最适合当前配置的高效算法,来达到优化运行效率的问题
2.定义文件路径
#----------------------------------2.定义文件路径----------------------------------
"""
这部分会将当前项目添加到系统路径上,以使得项目中的模块可以调用。同时将相对路径保存到ROOT中,便于寻找项目中的文件
"""
FILE = Path(__file__).resolve() # __file__指的是当前文件(即detect.py),FILE最终保存着当前文件的绝对路径,比如D://yolov5/detect.py
ROOT = FILE.parents[0] # YOLOv5 root directory ROOT保存着当前项目的父目录,比如 D://yolov5
if str(ROOT) not in sys.path: # sys.path即当前python环境可以运行的路径,假如当前项目不在该路径中,就无法运行其中的模块,所以就需要加载路径
sys.path.append(str(ROOT)) # add ROOT to PATH 把ROOT添加到运行路径上
ROOT = Path(os.path.relpath(ROOT, Path.cwd())) # relative ROOT设置为相对路径
3.加载自定义模块
#----------------------------------3.加载自定义模块----------------------------------
from models.common import DetectMultiBackend # 定义了图像处理 NMS等函数
from utils.datasets import IMG_FORMATS, VID_FORMATS, LoadImages, LoadStreams
# 定义了loadImages和LoadSteams 可以加载图像或视频并进行一些预处理
from utils.general import (LOGGER, check_file, check_img_size, check_imshow, check_requirements, colorstr,
increment_path, non_max_suppression, print_args, scale_coords, strip_optimizer, xyxy2xywh)
# 定义了常用工具函数,如检测文件是否存在,图像大小是否复合要求等
from utils.plots import Annotator, colors, save_one_box# 定义了Annotator 可以在图像上绘制矩形框和标注信息
from utils.torch_utils import select_device, time_sync# 定义了pytorch有关工具函数
4.run
#----------------------------------4.载入参数----------------------------------
@torch.no_grad()
def run(weights=ROOT / 'yolov5s.pt', # model.pt path(s)
source=ROOT / 'data/images', # file/dir/URL/glob, 0 for webcam
imgsz=640, # inference size (pixels)
conf_thres=0.25, # confidence threshold
iou_thres=0.45, # NMS IOU threshold
max_det=1000, # maximum detections per image
device='', # cuda device, i.e. 0 or 0,1,2,3 or cpu
view_img=False, # show results
save_txt=False, # save results to *.txt
save_conf=False, # save confidences in --save-txt labels
save_crop=False, # save cropped prediction boxes
nosave=False, # do not save images/videos
classes=None, # filter by class: --class 0, or --class 0 2 3
agnostic_nms=False, # class-agnostic NMS
augment=False, # augmented inference
visualize=False, # visualize features
update=False, # update all models
project=ROOT / 'runs/detect', # save results to project/name
name='exp', # save results to project/name
exist_ok=False, # existing project/name ok, do not increment
line_thickness=3, # bounding box thickness (pixels)
hide_labels=False, # hide labels
hide_conf=False, # hide confidences
half=False, # use FP16 half-precision inference
dnn=False, # use OpenCV DNN for ONNX inference
):
4.1参数初始化
# ----------------------------------4.1参数初始化----------------------------------
source = str(source)# 输入的路径变为字符串
save_img = not nosave and not source.endswith('.txt')# 是否保存图片和txt文件,如果nosave(传入的参数)为false且source的结尾不是txt则保存图片
# 判断source是不是视频/图像文件路径
# Path()提取文件名。suffix:最后一个组件的文件扩展名。若source是"YOLOv5/data/1.jpg", 则Path(source).suffix是".jpg", Path(source).suffix[1:]是"jpg"
# 而IMG_FORMATS 和 VID_FORMATS两个变量保存的是所有的视频和图片的格式后缀。
is_file = Path(source).suffix[1:] in (IMG_FORMATS + VID_FORMATS)
# 判断source是否是链接
# .lower()转化成小写 .upper()转化成大写 .title()首字符转化成大写,其余为小写, .startswith('http://')返回True or Flase
is_url = source.lower().startswith(('rtsp://', 'rtmp://', 'http://', 'https://'))
# 判断是source是否是摄像头
# .isnumeric()是否是由数字组成,返回True or False
webcam = source.isnumeric() or source.endswith('.txt') or (is_url and not is_file)
if is_url and is_file:
# 返回文件。如果source是一个指向图片/视频的链接,则下载输入数据
source = check_file(source) # download
4.2保存结果
# ----------------------------------4.2保存结果----------------------------------
# Directories
# save_dir是保存运行结果的文件夹名,是通过递增的方式来命名的。第一次运行时路径是“runs\detect\exp”,第二次运行时路径是“runs\detect\exp1”
save_dir = increment_path(Path(project) / name, exist_ok=exist_ok)
# 根据前面生成的路径创建文件夹
(save_dir / 'labels' if save_txt else save_dir).mkdir(parents=True, exist_ok=True) # make dir
4.3加载模型
# ----------------------------------4.3加载模型----------------------------------
# Load model 加载模型
"""
主要是选择谁不,初始化模型和检查图像大小
"""
device = select_device(device)# 获取设备 CPU/CUDA
model = DetectMultiBackend(weights, device=device, dnn=dnn)# DetectMultiBackend定义在models.common模块中,是我们要加载的网络,其中weights参数就是输入时指定的权重文件(比如yolov5s.pt)
"""
stride:推理时所用到的步长,默认为32, 大步长适合于大目标,小步长适合于小目标
names:保存推理结果名的列表,比如默认模型的值是['person', 'bicycle', 'car', ...]
pt: 加载的是否是pytorch模型(也就是pt格式的文件)
jit:当某段代码即将第一次被执行时进行编译,因而叫“即时编译”
onnx:利用Pytorch我们可以将model.pt转化为model.onnx格式的权重,在这里onnx充当一个后缀名称,
model.onnx就代表ONNX格式的权重文件,这个权重文件不仅包含了权重值,也包含了神经网络的网络流动信息以及每一层网络的输入输出信息和一些其他的辅助信息。
"""
stride, names, pt, jit, onnx = model.stride, model.names, model.pt, model.jit, model.onnx
imgsz = check_img_size(imgsz, s=stride) # 确保输入图片的尺寸imgsz能整除stride=32 如果不能则调整为能被整除并返回
# Half
half &= pt and device.type != 'cpu' # 如果不是CPU,使用半进度(图片半精度/模型半精度)
if pt:
model.model.half() if half else model.model.float()
4.4加载数据
# ----------------------------------4.4加载数据----------------------------------
# Dataloader
# 通过不同的输入源来设置不同的数据加载方式
if webcam: # 使用摄像头作为输入
view_img = check_imshow() # 检测cv2.imshow()方法是否可以执行,不能执行则抛出异常
cudnn.benchmark = True # set True to speed up constant image size inference 该设置可以加速预测
dataset = LoadStreams(source, img_size=imgsz, stride=stride, auto=pt and not jit) # 加载输入数据流
"""
source:输入数据源;image_size 图片识别前被放缩的大小;stride:识别时的步长,
auto的作用可以看utils.augmentations.letterbox方法,它决定了是否需要将图片填充为正方形,如果auto=True则不需要
"""
bs = len(dataset) # batch_size 批大小
else: # 直接从source文件下读取图片
dataset = LoadImages(source, img_size=imgsz, stride=stride, auto=pt and not jit)
bs = 1 # batch_size
# 保存视频的路径
vid_path, vid_writer = [None] * bs, [None] * bs # 前者是视频路径,后者是一个cv2.VideoWriter对象
4.5推理热身
# ----------------------------------4.5推理热身----------------------------------
"""
对模型机械能一些预处理以加速后续推理过程
1.将图片转化为Tensor格式 并根据要求转化为FP16或FP32
2.将像素从0-255归一化,并为批处理增加一维
3.激素时间消耗并更新dt列表
"""
# Run inference
if pt and device.type != 'cpu':
# 使用空白图片(零矩阵)预先用GPU跑一遍预测流程,可以加速预测
model(torch.zeros(1, 3, *imgsz).to(device).type_as(next(model.model.parameters()))) # warmup
dt, seen = [0.0, 0.0, 0.0], 0
'''
dt: 存储每一步骤的耗时
seen: 计数功能,已经处理完了多少帧图片
'''
# 去遍历图片,进行计数,
for path, im, im0s, vid_cap, s in dataset:
'''
在dataset中,每次迭代的返回值是self.sources, img, img0, None, ''
path:文件路径(即source)
im: resize后的图片(经过了放缩操作)
im0s: 原始图片
vid_cap=none
s: 图片的基本信息,比如路径,大小
'''
# ===以下部分是做预处理===#
t1 = time_sync() # 获取当前时间
im = torch.from_numpy(im).to(device) # 将图片放到指定设备(如GPU)上识别。#torch.size=[3,640,480]
im = im.half() if half else im.float() # uint8 to fp16/32 # 把输入从整型转化为半精度/全精度浮点数。
im /= 255 # 0 - 255 to 0.0 - 1.0 归一化,所有像素点除以255
if len(im.shape) == 3:
im = im[None] # expand for batch dim 添加一个第0维。缺少batch这个尺寸,所以将它扩充一下,变成[1,3,640,480]
t2 = time_sync() # 获取当前时间
dt[0] += t2 - t1 # 记录该阶段耗时
4.6前向推理
# ----------------------------------4.6前向推理----------------------------------
# Inference
# 可视化文件路径。如果为True则保留推理过程中的特征图,保存在runs文件夹中
visualize = increment_path(save_dir / Path(path).stem, mkdir=True) if visualize else False
# 推理结果,使用model函数对图像im进行预测,参数用于指示是否在预测时使用数据增强和可视化
pred = model(im, augment=augment, visualize=visualize) # 模型预测出来的所有检测框,torch.size=[1,18900,85]
t3 = time_sync()# 记录当前时间
dt[1] += t3 - t2
4.7非极大值抑制
# ----------------------------------4.7非极大值抑制NMS去掉多余anchor----------------------------------
# NMS
# 执行非极大值抑制,返回值为过滤后的预测框
pred = non_max_suppression(pred, conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det)
"""
pred: 网络的输出结果
conf_thres: 置信度阈值
iou_thres: iou阈值
classes: 是否只保留特定的类别 默认为None
agnostic_nms: 进行nms是否也去除不同类别之间的框
max_det: 检测框结果的最大数量 默认1000
"""
dt[2] += time_sync() - t3# 预测+NMS的时间
# Second-stage classifier (optional)
# pred = utils.general.apply_classifier(pred, classifier_model, im, im0s)
4.8预测过程
# ----------------------------------4.8预测过程----------------------------------
# Process predictions
# 把所有的检测框画到原图中
for i, det in enumerate(pred): # per image 每次迭代处理一张图片
"""
i:每个batch的信息
det:表示5个检测框的信息
"""
seen += 1 # seen是一个计数的功能
if webcam: # batch_size >= 1
# 如果输入源是webcam则batch_size>=1 取出dataset中的一张图片
p, im0, frame = path[i], im0s[i].copy(), dataset.count
s += f'{i}: ' # s后面拼接一个字符串i
else:
p, im0, frame = path, im0s.copy(), getattr(dataset, 'frame', 0)
"""
大部分我们一般都是从LoadImages流读取本都文件中的照片或者视频 所以batch_size=1
p: 当前图片/视频的绝对路径 如 F:\yolo_v5\yolov5-U\data\images\bus.jpg
s: 输出信息 初始为 ''
im0: 原始图片 letterbox + pad 之前的图片
frame: 视频流,此次取的是第几张图片
"""
p = Path(p) # to Path
# 图片/视频的保存路径save_path 如 runs\\detect\\exp8\\fire.jpg
save_path = str(save_dir / p.name) # im.jpg
# 设置保存框坐标的txt文件路径,每张图片对应一个框坐标信息
txt_path = str(save_dir / 'labels' / p.stem) + ('' if dataset.mode == 'image' else f'_{frame}') # im.txt
# 设置输出图片信息。图片shape (w, h)
s += '%gx%g ' % im.shape[2:] # print string
# 得到原图的宽和高
gn = torch.tensor(im0.shape)[[1, 0, 1, 0]] # normalization gain whwh
# 保存截图。如果save_crop的值为true,则将检测到的bounding_box单独保存成一张图片。
imc = im0.copy() if save_crop else im0 # for save_crop
# 得到一个绘图的类,类中预先存储了原图、线条宽度、类名
annotator = Annotator(im0, line_width=line_thickness, example=str(names))
# 判断有没有框
if len(det):
# Rescale boxes from img_size to im0 size
# 将预测信息映射到原图
# 将标注的bounding_box大小调整为和原图一致(因为训练时原图经过了放缩)此时坐标格式为xyxy
det[:, :4] = scale_coords(im.shape[2:], det[:, :4], im0.shape).round() # scale_coords:坐标映射功能
# 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
4.9打印目标检测结果
# ----------------------------------4.9打印目标检测结果----------------------------------
# Write results
# 保存预测结果:txt/图片画框/crop-image
for *xyxy, conf, cls in reversed(det):
# 将每个图片的预测信息分别存入save_dir/labels下的xxx.txt中 每行: class_id + score + xywh
if save_txt: # Write to file 保存txt文件
# 将xyxy(左上角+右下角)格式转为xywh(中心点+宽长)格式,并归一化,转化为列表再保存
xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist() # normalized xywh
# line的形式是: ”类别 x y w h“,若save_conf为true,则line的形式是:”类别 x y w h 置信度“
line = (cls, *xywh, conf) if save_conf else (cls, *xywh) # label format
with open(txt_path + '.txt', 'a') as f:
# 写入对应的文件夹里,路径默认为“runs\detect\exp*\labels”
f.write(('%g ' * len(line)).rstrip() % line + '\n')
# 在原图上画框+将预测到的目标剪切出来保存成图片,保存在save_dir/crops下,在原图像画图或者保存结果
if save_img or save_crop or view_img: # Add bbox to image
c = int(cls) # integer class # 类别标号
label = None if hide_labels else (
names[c] if hide_conf else f'{names[c]} {conf:.2f}') # 类别名
annotator.box_label(xyxy, label, color=colors(c, True)) # 绘制边框
# 在原图上画框+将预测到的目标剪切出来保存成图片,保存在save_dir/crops下(单独保存)
if save_crop:
save_one_box(xyxy, imc, file=save_dir / 'crops' / names[c] / f'{p.stem}.jpg', BGR=True)
4.10查看实时检测结果
# ----------------------------------4.10查看实时检测结果----------------------------------
# Print time (inference-only)
# 打印耗时
LOGGER.info(f'{s}Done. ({t3 - t2:.3f}s)')
# Stream results
# 如果设置展示,则show图片 / 视频
im0 = annotator.result() # im0是绘制好的图片
# 显示图片
if view_img:
cv2.imshow(str(p), im0)
cv2.waitKey(1) # 暂停 1 millisecond
4.11设置保存结果
# ----------------------------------4.11设置保存结果----------------------------------
# Save results (image with detections)
if save_img:
if dataset.mode == 'image': # 保存图片
cv2.imwrite(save_path, im0)
else: # 'video' or 'stream' #保存视频
if vid_path[i] != save_path: # new video
vid_path[i] = save_path
if isinstance(vid_writer[i], cv2.VideoWriter):
vid_writer[i].release() # release previous video writer
if vid_cap: # video
fps = vid_cap.get(cv2.CAP_PROP_FPS)# 视频FPS
w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))# 获取视频帧宽度
h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))# 获取视频帧高度
else: # stream
fps, w, h = 30, im0.shape[1], im0.shape[0]
save_path += '.mp4'
vid_writer[i] = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
vid_writer[i].write(im0)
4.12终端打印结果
# ----------------------------------4.12终端打印结果----------------------------------
# Print results
t = tuple(x / seen * 1E3 for x in dt) # speeds per image # 平均每张图片所耗时间
LOGGER.info(f'Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {(1, 3, *imgsz)}' % t)
if save_txt or save_img:
s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else ''
LOGGER.info(f"Results saved to {colorstr('bold', save_dir)}{s}")# 标签保存路径
if update:
strip_optimizer(weights) # update model (to fix SourceChangeWarning)
5.设置OTP参数
# ----------------------------------5.设置OTP参数----------------------------------
def parse_opt():
parser = argparse.ArgumentParser()
# weights:训练的权重路径
parser.add_argument('--weights', nargs='+', type=str, default=ROOT / 'yolov5s.pt', help='model path(s)')
# source:测试数据 0是电脑摄像头
parser.add_argument('--source', type=str, default=ROOT / 'data/images', help='file/dir/URL/glob, 0 for webcam')
# imgsz:预测时网络输入图片的尺寸
parser.add_argument('--imgsz', '--img', '--img-size', nargs='+', type=int, default=[640], help='inference size h,w')
# conf-thres:置信度阈值
parser.add_argument('--conf-thres', type=float, default=0.5, help='confidence threshold')
# iou-thres:非极大抑制时的loU阈值
parser.add_argument('--iou-thres', type=float, default=0.45, help='NMS IoU threshold')
# max-det:保留的最大检测框数量,每张图片中检测目标的个数最多为1000类·
parser.add_argument('--max-det', type=int, default=1000, help='maximum detections per image')
# device:使用的设备,可以是cuda 设备的ID
parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
# view-img: 是否展示预测之后的图片/视频
parser.add_argument('--view-img', action='store_true', help='show results')
# save-txt:是否anhor框坐标以txt保存
parser.add_argument('--save-txt', action='store_true', help='save results to *.txt')
# save-conf:预测的txt文件是否保存检测结果的置信度到txt文件
parser.add_argument('--save-conf', action='store_true', help='save confidences in --save-txt labels')
# save-crop:是否保存裁剪预测图片
parser.add_argument('--save-crop', action='store_true', help='save cropped prediction boxes')
# nosave:不保存图片、视频,要保存图片
parser.add_argument('--nosave', action='store_true', help='do not save images/videos')
# classes:仅检测指定类别
parser.add_argument('--classes', nargs='+', type=int, help='filter by class: --classes 0, or --classes 0 2 3')
# agnostic-nms:是否使用类别不敏感的非极大抑制
parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS')
# augment:是否使用数据增强进行推理
parser.add_argument('--augment', action='store_true', help='augmented inference')
# visualize:是否可视化特征图
parser.add_argument('--visualize', action='store_true', help='visualize features')
# update:如果为True,则对所有模型进行strip_optimizer操作,去除pt文件中的优化器等信息,默认为False
parser.add_argument('--update', action='store_true', help='update all models')
# project:结果保存的项目目录路径,默认为'ROOT/runs/detect'
parser.add_argument('--project', default=ROOT / 'runs/detect', help='save results to project/name')
# name:结果保存的子目录名称
parser.add_argument('--name', default='exp', help='save results to project/name')
# exist-ok:是否覆盖已有结果
parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
# line-thickness:画bounding box时的线条宽度,默认为3
parser.add_argument('--line-thickness', default=3, type=int, help='bounding box thickness (pixels)')
# hide-labels:是否隐藏标签信息
parser.add_argument('--hide-labels', default=False, action='store_true', help='hide labels')
# hide-conf:是否隐藏置信度信息
parser.add_argument('--hide-conf', default=False, action='store_true', help='hide confidences')
# half: 是否使用FP16半精度进行推理
parser.add_argument('--half', action='store_true', help='use FP16 half-precision inference')
# dnn:是否使用OpenCV DNN进行ONNX推理
parser.add_argument('--dnn', action='store_true', help='use OpenCV DNN for ONNX inference')
opt = parser.parse_args()# 扩充维度
opt.imgsz *= 2 if len(opt.imgsz) == 1 else 1 # expand
print_args(FILE.stem, opt)# 打印所有参数信息
return opt
6.main函数
# ----------------------------------6.main函数----------------------------------
def main(opt):
check_requirements(exclude=('tensorboard', 'thop'))# 检查环境,主要是requirements.txt是否安装
run(**vars(opt))# 运行run函数
# 命令使用
# python detect.py --weights runs/train/exp_yolov5s/weights/best.pt --source data/images/fishman.jpg # webcam
if __name__ == "__main__":
opt = parse_opt()# 解析opt参数
main(opt)# 执行主函数
7.整体代码
# YOLOv5 🚀 by Ultralytics, GPL-3.0 license
"""
Run inference on images, videos, directories, streams, etc.
Usage:
$ python path/to/detect.py --weights yolov5s.pt --source 0 # webcam
img.jpg # image
vid.mp4 # video
path/ # directory
path/*.jpg # glob
'https://youtu.be/Zgi9g1ksQHc' # YouTube
'rtsp://example.com/media.mp4' # RTSP, RTMP, HTTP stream
"""
#----------------------------------1.导入Python库----------------------------------
"""它是一个用于命令项选项与参数解析的模块,通过在程序中定义好我们需要的参数,argparse将会从sys.argv中解析出这些参数,并自动生成帮助和使用信息"""
import argparse # 解析命令行参数的库
"""它提供了多种操作系统的接口。通过os模块提供的操作系统接口,我们可以对操作系统里文件、终端、进程等进行操作"""
import os # 与操作系统进行交互的文件库 包含文件路径操作与解析
"""它是与python解释器交互的一个接口,该模块提供对解释器使用或维护的一些变量的访问和获取,它提供了许多函数和变量来处理Python 运行时环境的不同部分"""
import sys # sys模块包含了与python解释器和它的环境有关的函数。
"""这个库提供了一种面向对象的方式来与文件系统交互,可以让代码更简洁、更易读"""
from pathlib import Path # Path能够更加方便得对字符串路径进行处理
"""opencv库"""
import cv2 # opencv库
"""这是主要的Pytorch库。它提供了构建、训练和评估神经网络的工具"""
import torch # pytorch 深度学习库
"""它提供了一个接口,用于使用cuDNN库,在NVIDIA GPU上高效地进行深度学习。cudnn模块是一个Pytorch库的扩展"""
import torch.backends.cudnn as cudnn # 让内置的cudnn的 auto-tuner 自动寻找最适合当前配置的高效算法,来达到优化运行效率的问题
#----------------------------------2.定义文件路径----------------------------------
"""
这部分会将当前项目添加到系统路径上,以使得项目中的模块可以调用。同时将相对路径保存到ROOT中,便于寻找项目中的文件
"""
FILE = Path(__file__).resolve() # __file__指的是当前文件(即detect.py),FILE最终保存着当前文件的绝对路径,比如D://yolov5/detect.py
ROOT = FILE.parents[0] # YOLOv5 root directory ROOT保存着当前项目的父目录,比如 D://yolov5
if str(ROOT) not in sys.path: # sys.path即当前python环境可以运行的路径,假如当前项目不在该路径中,就无法运行其中的模块,所以就需要加载路径
sys.path.append(str(ROOT)) # add ROOT to PATH 把ROOT添加到运行路径上
ROOT = Path(os.path.relpath(ROOT, Path.cwd())) # relative ROOT设置为相对路径
#----------------------------------3.加载自定义模块----------------------------------
from models.common import DetectMultiBackend # 定义了图像处理 NMS等函数
from utils.datasets import IMG_FORMATS, VID_FORMATS, LoadImages, LoadStreams
# 定义了loadImages和LoadSteams 可以加载图像或视频并进行一些预处理
from utils.general import (LOGGER, check_file, check_img_size, check_imshow, check_requirements, colorstr,
increment_path, non_max_suppression, print_args, scale_coords, strip_optimizer, xyxy2xywh)
# 定义了常用工具函数,如检测文件是否存在,图像大小是否复合要求等
from utils.plots import Annotator, colors, save_one_box# 定义了Annotator 可以在图像上绘制矩形框和标注信息
from utils.torch_utils import select_device, time_sync# 定义了pytorch有关工具函数
#----------------------------------4.载入参数----------------------------------
@torch.no_grad()
def run(weights=ROOT / 'yolov5s.pt', # model.pt path(s)
source=ROOT / 'data/images', # file/dir/URL/glob, 0 for webcam
imgsz=640, # inference size (pixels)
conf_thres=0.25, # confidence threshold
iou_thres=0.45, # NMS IOU threshold
max_det=1000, # maximum detections per image
device='', # cuda device, i.e. 0 or 0,1,2,3 or cpu
view_img=False, # show results
save_txt=False, # save results to *.txt
save_conf=False, # save confidences in --save-txt labels
save_crop=False, # save cropped prediction boxes
nosave=False, # do not save images/videos
classes=None, # filter by class: --class 0, or --class 0 2 3
agnostic_nms=False, # class-agnostic NMS
augment=False, # augmented inference
visualize=False, # visualize features
update=False, # update all models
project=ROOT / 'runs/detect', # save results to project/name
name='exp', # save results to project/name
exist_ok=False, # existing project/name ok, do not increment
line_thickness=3, # bounding box thickness (pixels)
hide_labels=False, # hide labels
hide_conf=False, # hide confidences
half=False, # use FP16 half-precision inference
dnn=False, # use OpenCV DNN for ONNX inference
):
# ----------------------------------4.1参数初始化----------------------------------
source = str(source)# 输入的路径变为字符串
save_img = not nosave and not source.endswith('.txt')# 是否保存图片和txt文件,如果nosave(传入的参数)为false且source的结尾不是txt则保存图片
# 判断source是不是视频/图像文件路径
# Path()提取文件名。suffix:最后一个组件的文件扩展名。若source是"YOLOv5/data/1.jpg", 则Path(source).suffix是".jpg", Path(source).suffix[1:]是"jpg"
# 而IMG_FORMATS 和 VID_FORMATS两个变量保存的是所有的视频和图片的格式后缀。
is_file = Path(source).suffix[1:] in (IMG_FORMATS + VID_FORMATS)
# 判断source是否是链接
# .lower()转化成小写 .upper()转化成大写 .title()首字符转化成大写,其余为小写, .startswith('http://')返回True or Flase
is_url = source.lower().startswith(('rtsp://', 'rtmp://', 'http://', 'https://'))
# 判断是source是否是摄像头
# .isnumeric()是否是由数字组成,返回True or False
webcam = source.isnumeric() or source.endswith('.txt') or (is_url and not is_file)
if is_url and is_file:
# 返回文件。如果source是一个指向图片/视频的链接,则下载输入数据
source = check_file(source) # download
# ----------------------------------4.2保存结果----------------------------------
# Directories
# save_dir是保存运行结果的文件夹名,是通过递增的方式来命名的。第一次运行时路径是“runs\detect\exp”,第二次运行时路径是“runs\detect\exp1”
save_dir = increment_path(Path(project) / name, exist_ok=exist_ok)
# 根据前面生成的路径创建文件夹
(save_dir / 'labels' if save_txt else save_dir).mkdir(parents=True, exist_ok=True) # make dir
# ----------------------------------4.3加载模型----------------------------------
# Load model 加载模型
"""
主要是选择谁不,初始化模型和检查图像大小
"""
device = select_device(device)# 获取设备 CPU/CUDA
model = DetectMultiBackend(weights, device=device, dnn=dnn)# DetectMultiBackend定义在models.common模块中,是我们要加载的网络,其中weights参数就是输入时指定的权重文件(比如yolov5s.pt)
"""
stride:推理时所用到的步长,默认为32, 大步长适合于大目标,小步长适合于小目标
names:保存推理结果名的列表,比如默认模型的值是['person', 'bicycle', 'car', ...]
pt: 加载的是否是pytorch模型(也就是pt格式的文件)
jit:当某段代码即将第一次被执行时进行编译,因而叫“即时编译”
onnx:利用Pytorch我们可以将model.pt转化为model.onnx格式的权重,在这里onnx充当一个后缀名称,
model.onnx就代表ONNX格式的权重文件,这个权重文件不仅包含了权重值,也包含了神经网络的网络流动信息以及每一层网络的输入输出信息和一些其他的辅助信息。
"""
stride, names, pt, jit, onnx = model.stride, model.names, model.pt, model.jit, model.onnx
imgsz = check_img_size(imgsz, s=stride) # 确保输入图片的尺寸imgsz能整除stride=32 如果不能则调整为能被整除并返回
# Half
half &= pt and device.type != 'cpu' # 如果不是CPU,使用半进度(图片半精度/模型半精度)
if pt:
model.model.half() if half else model.model.float()
# ----------------------------------4.4加载数据----------------------------------
# Dataloader
# 通过不同的输入源来设置不同的数据加载方式
if webcam: # 使用摄像头作为输入
view_img = check_imshow() # 检测cv2.imshow()方法是否可以执行,不能执行则抛出异常
cudnn.benchmark = True # set True to speed up constant image size inference 该设置可以加速预测
dataset = LoadStreams(source, img_size=imgsz, stride=stride, auto=pt and not jit) # 加载输入数据流
"""
source:输入数据源;image_size 图片识别前被放缩的大小;stride:识别时的步长,
auto的作用可以看utils.augmentations.letterbox方法,它决定了是否需要将图片填充为正方形,如果auto=True则不需要
"""
bs = len(dataset) # batch_size 批大小
else: # 直接从source文件下读取图片
dataset = LoadImages(source, img_size=imgsz, stride=stride, auto=pt and not jit)
bs = 1 # batch_size
# 保存视频的路径
vid_path, vid_writer = [None] * bs, [None] * bs # 前者是视频路径,后者是一个cv2.VideoWriter对象
# ----------------------------------4.5推理热身----------------------------------
"""
对模型机械能一些预处理以加速后续推理过程
1.将图片转化为Tensor格式 并根据要求转化为FP16或FP32
2.将像素从0-255归一化,并为批处理增加一维
3.激素时间消耗并更新dt列表
"""
# Run inference
if pt and device.type != 'cpu':
# 使用空白图片(零矩阵)预先用GPU跑一遍预测流程,可以加速预测
model(torch.zeros(1, 3, *imgsz).to(device).type_as(next(model.model.parameters()))) # warmup
dt, seen = [0.0, 0.0, 0.0], 0
'''
dt: 存储每一步骤的耗时
seen: 计数功能,已经处理完了多少帧图片
'''
# 去遍历图片,进行计数,
for path, im, im0s, vid_cap, s in dataset:
'''
在dataset中,每次迭代的返回值是self.sources, img, img0, None, ''
path:文件路径(即source)
im: resize后的图片(经过了放缩操作)
im0s: 原始图片
vid_cap=none
s: 图片的基本信息,比如路径,大小
'''
# ===以下部分是做预处理===#
t1 = time_sync() # 获取当前时间
im = torch.from_numpy(im).to(device) # 将图片放到指定设备(如GPU)上识别。#torch.size=[3,640,480]
im = im.half() if half else im.float() # uint8 to fp16/32 # 把输入从整型转化为半精度/全精度浮点数。
im /= 255 # 0 - 255 to 0.0 - 1.0 归一化,所有像素点除以255
if len(im.shape) == 3:
im = im[None] # expand for batch dim 添加一个第0维。缺少batch这个尺寸,所以将它扩充一下,变成[1,3,640,480]
t2 = time_sync() # 获取当前时间
dt[0] += t2 - t1 # 记录该阶段耗时
# ----------------------------------4.6前向推理----------------------------------
# Inference
# 可视化文件路径。如果为True则保留推理过程中的特征图,保存在runs文件夹中
visualize = increment_path(save_dir / Path(path).stem, mkdir=True) if visualize else False
# 推理结果,使用model函数对图像im进行预测,参数用于指示是否在预测时使用数据增强和可视化
pred = model(im, augment=augment, visualize=visualize) # 模型预测出来的所有检测框,torch.size=[1,18900,85]
t3 = time_sync()# 记录当前时间
dt[1] += t3 - t2
# ----------------------------------4.7非极大值抑制NMS去掉多余anchor----------------------------------
# NMS
# 执行非极大值抑制,返回值为过滤后的预测框
pred = non_max_suppression(pred, conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det)
"""
pred: 网络的输出结果
conf_thres: 置信度阈值
iou_thres: iou阈值
classes: 是否只保留特定的类别 默认为None
agnostic_nms: 进行nms是否也去除不同类别之间的框
max_det: 检测框结果的最大数量 默认1000
"""
dt[2] += time_sync() - t3# 预测+NMS的时间
# Second-stage classifier (optional)
# pred = utils.general.apply_classifier(pred, classifier_model, im, im0s)
# ----------------------------------4.8预测过程----------------------------------
# Process predictions
# 把所有的检测框画到原图中
for i, det in enumerate(pred): # per image 每次迭代处理一张图片
"""
i:每个batch的信息
det:表示5个检测框的信息
"""
seen += 1 # seen是一个计数的功能
if webcam: # batch_size >= 1
# 如果输入源是webcam则batch_size>=1 取出dataset中的一张图片
p, im0, frame = path[i], im0s[i].copy(), dataset.count
s += f'{i}: ' # s后面拼接一个字符串i
else:
p, im0, frame = path, im0s.copy(), getattr(dataset, 'frame', 0)
"""
大部分我们一般都是从LoadImages流读取本都文件中的照片或者视频 所以batch_size=1
p: 当前图片/视频的绝对路径 如 F:\yolo_v5\yolov5-U\data\images\bus.jpg
s: 输出信息 初始为 ''
im0: 原始图片 letterbox + pad 之前的图片
frame: 视频流,此次取的是第几张图片
"""
p = Path(p) # to Path
# 图片/视频的保存路径save_path 如 runs\\detect\\exp8\\fire.jpg
save_path = str(save_dir / p.name) # im.jpg
# 设置保存框坐标的txt文件路径,每张图片对应一个框坐标信息
txt_path = str(save_dir / 'labels' / p.stem) + ('' if dataset.mode == 'image' else f'_{frame}') # im.txt
# 设置输出图片信息。图片shape (w, h)
s += '%gx%g ' % im.shape[2:] # print string
# 得到原图的宽和高
gn = torch.tensor(im0.shape)[[1, 0, 1, 0]] # normalization gain whwh
# 保存截图。如果save_crop的值为true,则将检测到的bounding_box单独保存成一张图片。
imc = im0.copy() if save_crop else im0 # for save_crop
# 得到一个绘图的类,类中预先存储了原图、线条宽度、类名
annotator = Annotator(im0, line_width=line_thickness, example=str(names))
# 判断有没有框
if len(det):
# Rescale boxes from img_size to im0 size
# 将预测信息映射到原图
# 将标注的bounding_box大小调整为和原图一致(因为训练时原图经过了放缩)此时坐标格式为xyxy
det[:, :4] = scale_coords(im.shape[2:], det[:, :4], im0.shape).round() # scale_coords:坐标映射功能
# 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
# ----------------------------------4.9打印目标检测结果----------------------------------
# Write results
# 保存预测结果:txt/图片画框/crop-image
for *xyxy, conf, cls in reversed(det):
# 将每个图片的预测信息分别存入save_dir/labels下的xxx.txt中 每行: class_id + score + xywh
if save_txt: # Write to file 保存txt文件
# 将xyxy(左上角+右下角)格式转为xywh(中心点+宽长)格式,并归一化,转化为列表再保存
xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist() # normalized xywh
# line的形式是: ”类别 x y w h“,若save_conf为true,则line的形式是:”类别 x y w h 置信度“
line = (cls, *xywh, conf) if save_conf else (cls, *xywh) # label format
with open(txt_path + '.txt', 'a') as f:
# 写入对应的文件夹里,路径默认为“runs\detect\exp*\labels”
f.write(('%g ' * len(line)).rstrip() % line + '\n')
# 在原图上画框+将预测到的目标剪切出来保存成图片,保存在save_dir/crops下,在原图像画图或者保存结果
if save_img or save_crop or view_img: # Add bbox to image
c = int(cls) # integer class # 类别标号
label = None if hide_labels else (
names[c] if hide_conf else f'{names[c]} {conf:.2f}') # 类别名
annotator.box_label(xyxy, label, color=colors(c, True)) # 绘制边框
# 在原图上画框+将预测到的目标剪切出来保存成图片,保存在save_dir/crops下(单独保存)
if save_crop:
save_one_box(xyxy, imc, file=save_dir / 'crops' / names[c] / f'{p.stem}.jpg', BGR=True)
# ----------------------------------4.10查看实时检测结果----------------------------------
# Print time (inference-only)
# 打印耗时
LOGGER.info(f'{s}Done. ({t3 - t2:.3f}s)')
# Stream results
# 如果设置展示,则show图片 / 视频
im0 = annotator.result() # im0是绘制好的图片
# 显示图片
if view_img:
cv2.imshow(str(p), im0)
cv2.waitKey(1) # 暂停 1 millisecond
# ----------------------------------4.11设置保存结果----------------------------------
# Save results (image with detections)
if save_img:
if dataset.mode == 'image': # 保存图片
cv2.imwrite(save_path, im0)
else: # 'video' or 'stream' #保存视频
if vid_path[i] != save_path: # new video
vid_path[i] = save_path
if isinstance(vid_writer[i], cv2.VideoWriter):
vid_writer[i].release() # release previous video writer
if vid_cap: # video
fps = vid_cap.get(cv2.CAP_PROP_FPS)# 视频FPS
w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))# 获取视频帧宽度
h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))# 获取视频帧高度
else: # stream
fps, w, h = 30, im0.shape[1], im0.shape[0]
save_path += '.mp4'
vid_writer[i] = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
vid_writer[i].write(im0)
# ----------------------------------4.12终端打印结果----------------------------------
# Print results
t = tuple(x / seen * 1E3 for x in dt) # speeds per image # 平均每张图片所耗时间
LOGGER.info(f'Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {(1, 3, *imgsz)}' % t)
if save_txt or save_img:
s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else ''
LOGGER.info(f"Results saved to {colorstr('bold', save_dir)}{s}")# 标签保存路径
if update:
strip_optimizer(weights) # update model (to fix SourceChangeWarning)
# ----------------------------------5.设置OTP参数----------------------------------
def parse_opt():
parser = argparse.ArgumentParser()
# weights:训练的权重路径
parser.add_argument('--weights', nargs='+', type=str, default=ROOT / 'yolov5s.pt', help='model path(s)')
# source:测试数据 0是电脑摄像头
parser.add_argument('--source', type=str, default=ROOT / 'data/images', help='file/dir/URL/glob, 0 for webcam')
# imgsz:预测时网络输入图片的尺寸
parser.add_argument('--imgsz', '--img', '--img-size', nargs='+', type=int, default=[640], help='inference size h,w')
# conf-thres:置信度阈值
parser.add_argument('--conf-thres', type=float, default=0.5, help='confidence threshold')
# iou-thres:非极大抑制时的loU阈值
parser.add_argument('--iou-thres', type=float, default=0.45, help='NMS IoU threshold')
# max-det:保留的最大检测框数量,每张图片中检测目标的个数最多为1000类·
parser.add_argument('--max-det', type=int, default=1000, help='maximum detections per image')
# device:使用的设备,可以是cuda 设备的ID
parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
# view-img: 是否展示预测之后的图片/视频
parser.add_argument('--view-img', action='store_true', help='show results')
# save-txt:是否anhor框坐标以txt保存
parser.add_argument('--save-txt', action='store_true', help='save results to *.txt')
# save-conf:预测的txt文件是否保存检测结果的置信度到txt文件
parser.add_argument('--save-conf', action='store_true', help='save confidences in --save-txt labels')
# save-crop:是否保存裁剪预测图片
parser.add_argument('--save-crop', action='store_true', help='save cropped prediction boxes')
# nosave:不保存图片、视频,要保存图片
parser.add_argument('--nosave', action='store_true', help='do not save images/videos')
# classes:仅检测指定类别
parser.add_argument('--classes', nargs='+', type=int, help='filter by class: --classes 0, or --classes 0 2 3')
# agnostic-nms:是否使用类别不敏感的非极大抑制
parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS')
# augment:是否使用数据增强进行推理
parser.add_argument('--augment', action='store_true', help='augmented inference')
# visualize:是否可视化特征图
parser.add_argument('--visualize', action='store_true', help='visualize features')
# update:如果为True,则对所有模型进行strip_optimizer操作,去除pt文件中的优化器等信息,默认为False
parser.add_argument('--update', action='store_true', help='update all models')
# project:结果保存的项目目录路径,默认为'ROOT/runs/detect'
parser.add_argument('--project', default=ROOT / 'runs/detect', help='save results to project/name')
# name:结果保存的子目录名称
parser.add_argument('--name', default='exp', help='save results to project/name')
# exist-ok:是否覆盖已有结果
parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
# line-thickness:画bounding box时的线条宽度,默认为3
parser.add_argument('--line-thickness', default=3, type=int, help='bounding box thickness (pixels)')
# hide-labels:是否隐藏标签信息
parser.add_argument('--hide-labels', default=False, action='store_true', help='hide labels')
# hide-conf:是否隐藏置信度信息
parser.add_argument('--hide-conf', default=False, action='store_true', help='hide confidences')
# half: 是否使用FP16半精度进行推理
parser.add_argument('--half', action='store_true', help='use FP16 half-precision inference')
# dnn:是否使用OpenCV DNN进行ONNX推理
parser.add_argument('--dnn', action='store_true', help='use OpenCV DNN for ONNX inference')
opt = parser.parse_args()# 扩充维度
opt.imgsz *= 2 if len(opt.imgsz) == 1 else 1 # expand
print_args(FILE.stem, opt)# 打印所有参数信息
return opt
# ----------------------------------6.main函数----------------------------------
def main(opt):
check_requirements(exclude=('tensorboard', 'thop'))# 检查环境,主要是requirements.txt是否安装
run(**vars(opt))# 运行run函数
# 命令使用
# python detect.py --weights runs/train/exp_yolov5s/weights/best.pt --source data/images/fishman.jpg # webcam
if __name__ == "__main__":
opt = parse_opt()# 解析opt参数
main(opt)# 执行主函数
学习完detect.py后,一般就能确定自己电脑可以运行yolov5了,接下来便是给图像打标签然后进行训练自己的模型了,所以接下来会学习train.py.如果本文对你有所帮助,请点个赞呀!