yolov5检测部分代码解析
yolov5源代码
：https://github.com/ultralytics/yolov5
设置目标检测的配置参数：
parser = argparse.ArgumentParser()
parser.add_argument('--weights', nargs='+', type=str, default='yolov5s.pt', help='model.pt path(s)')
parser.add_argument('--source', type=str, default='data/images', help='source') # file/folder, 0 for webcam
parser.add_argument('--img-size', type=int, default=640, help='inference size (pixels)')
parser.add_argument('--conf-thres', type=float, default=0.25, help='object confidence threshold')
parser.add_argument('--iou-thres', type=float, default=0.45, help='IOU threshold for NMS')
parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
parser.add_argument('--view-img', action='store_true', help='display results')
parser.add_argument('--save-txt', action='store_true', help='save results to *.txt')
parser.add_argument('--save-conf', action='store_true', help='save confidences in --save-txt labels')
parser.add_argument('--nosave', action='store_true', help='do not save images/videos')
parser.add_argument('--classes', nargs='+', type=int, help='filter by class: --class 0, or --class 0 2 3')
parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS')
parser.add_argument('--augment', action='store_true', help='augmented inference')
parser.add_argument('--update', action='store_true', help='update all models')
parser.add_argument('--project', default='runs/detect', help='save results to project/name')
parser.add_argument('--name', default='exp', help='save results to project/name')
parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
opt = parser.parse_args()
print(opt)
执行检测程序前，查看requirements.txt里的相关安装包是否安装在执行检测的python环境下：
check_requirements(exclude=('pycocotools', 'thop'))
def check_requirements(file='requirements.txt', exclude=()):
# Check installed dependencies meet requirements
import pkg_resources as pkg
prefix = colorstr('red', 'bold', 'requirements:')
file = Path(file)
if not file.exists():
print(f"{prefix} {file.resolve()} not found, check failed.")
return
n = 0 # number of packages updates
requirements = [f'{x.name}{x.specifier}' for x in pkg.parse_requirements(file.open()) if x.name not in exclude]
for r in requirements:
try:
pkg.require(r)
except Exception as e: # DistributionNotFound or VersionConflict if requirements not met
n += 1
print(f"{prefix} {e.req} not found and is required by YOLOv5, attempting auto-update...")
print(subprocess.check_output(f"pip install '{e.req}'", shell=True).decode())
if n: # if packages updated
s = f"{prefix} {n} package{'s' * (n > 1)} updated per {file.resolve()}\n" \
f"{prefix} ⚠️ {colorstr('bold', 'Restart runtime or rerun command for updates to take effect')}\n"
print(emojis(s)) # emoji-safe
执行检测部分：
with torch.no_grad():
if opt.update: # update all models (to fix SourceChangeWarning)
for opt.weights in ['yolov5s.pt', 'yolov5m.pt', 'yolov5l.pt', 'yolov5x.pt']:
detect()
strip_optimizer(opt.weights)
else:
detect()
接下来主要看detect()内部的执行。
第一行是相关参数的加载，参数加载的内容为输入的配置参数。
第二行为是否保存图片的判断（bool类型），默认为保存图片，不保存图片可以在执行程序中输入 –nosave，或者你检测的对象是一个网络视频流的地址，其中视频流地址保存在一个文件类型为txt的文档里。
第三个是摄像头检测的一个判断（bool类型）， 是你在–source后的输入，当检测对象是本地的一个摄像头或者是一个txt文档内有网络视频源地址，或者是你直接输入网络视频源的地址以(‘rtsp://’, ‘rtmp://’, ‘http://’)这三个开头的网络地址，以上三种情况（不是指网络地址的三种情况）判定你的检测对象是摄像头检测。
def detect(save_img=False):
source, weights, view_img, save_txt, imgsz = opt.source, opt.weights, opt.view_img, opt.save_txt, opt.img_size
save_img = not opt.nosave and not source.endswith('.txt') # save inference images
webcam = source.isnumeric() or source.endswith('.txt') or source.lower().startswith(
('rtsp://', 'rtmp://', 'http://'))
保存检测结果地址的一个写入程序：
第一行为保存地址返回（str），其中opt.project，默认为’runs/detect’， opt.name默认为’exp’，opt.exist_ok默认为False，opt.exist_ok判断是否将检测结果保存在你默认或者输入的opt.project/opt.name文件内，执行检测程序时如果不改变opt.project和opt.name，只输入–exist_ok，则保存结果会在’runs/detect/exp文件夹内。如果不输入默认的opt.project，opt.name和opt.exist_ok ，increment_path()这个程序是通过通过正则匹配的方式在你的runs/detect文件内查找含有exp的文件，并匹配出最大的数字，进行+1后，返回’runs/detect/exp’ + ‘(maxnum+1)’的一个字符串
第二行为第一行返回保存位置的一个文件夹创建
# Directories
save_dir = Path(increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok)) # increment run
(save_dir / 'labels' if save_txt else save_dir).mkdir(parents=True, exist_ok=True) # make dir
def increment_path(path, exist_ok=True, sep=''):
# Increment path, i.e. runs/exp --> runs/exp{sep}0, runs/exp{sep}1 etc.
path = Path(path) # os-agnostic
if (path.exists() and exist_ok) or (not path.exists()):
return str(path)
else:
dirs = glob.glob(f"{path}{sep}*") # similar paths
matches = [re.search(rf"%s{sep}(\d+)" % path.stem, d) for d in dirs]
i = [int(m.groups()[0]) for m in matches if m] # indices
n = max(i) + 1 if i else 2 # increment number
return f"{path}{sep}{n}" # update path
第一行打印日志信息
第二行选择程序在哪个硬件执行（CPU或者GPU），select_device()，如果在opt.device选择输入cpu，则在cpu上进行检测程序的执行，如果不输入，则通过torch.cuda.is_available()来判断是否使用GPU运行程序
第三行通过降低少许检测精度提高检测速度的判断（bool）
# Initialize
set_logging()
device = select_device(opt.device)
half = device.type != 'cpu' # half precision only supported on CUDA
def select_device(device='', batch_size=None):
# device = 'cpu' or '0' or '0,1,2,3'
s = f'YOLOv5 🚀 {git_describe() or date_modified()} torch {torch.__version__} ' # string
cpu = device.lower() == 'cpu'
if cpu:
os.environ['CUDA_VISIBLE_DEVICES'] = '-1' # force torch.cuda.is_available() = False
elif device: # non-cpu device requested
os.environ['CUDA_VISIBLE_DEVICES'] = device # set environment variable
assert torch.cuda.is_available(), f'CUDA unavailable, invalid device {device} requested' # check availability
cuda = not cpu and torch.cuda.is_available()
if cuda:
n = torch.cuda.device_count()
if n > 1 and batch_size: # check that batch_size is compatible with device_count
assert batch_size % n == 0, f'batch-size {batch_size} not multiple of GPU count {n}'
space = ' ' * len(s)
for i, d in enumerate(device.split(',') if device else range(n)):
p = torch.cuda.get_device_properties(i)
s += f"{'' if i == 0 else space}CUDA:{d} ({p.name}, {p.total_memory / 1024 ** 2}MB)\n" # bytes to MB
else:
s += 'CPU\n'
logger.info(s.encode().decode('ascii', 'ignore') if platform.system() == 'Windows' else s) # emoji-safe
return torch.device('cuda:0' if cuda else 'cpu')
第一行为模型的载入，其中attempt_load()是对权重文件的一个载入，主要是通过torch.load()这个函数，值得注意的是，这个可以是多个模型的混和载入（需要注意每个模型的输入和输入大小），也可以是对一个完整检测模型的载入。
第二行是检测模型时yolov5也用了FPN里面的一个思路，将卷积神经网络中最后三层通过上采样的方式进行了一个堆叠的，要求图片的分辨率为32的一个整数倍。
第三行为输入判断你输入的opt.img_size是否是32的倍数，如果不是，则改变你输入的img_size为向上取32的倍数的分辨率。
# Load model
model = attempt_load(weights, map_location=device) # load FP32 model
stride = int(model.stride.max()) # model stride
imgsz = check_img_size(imgsz, s=stride) # check img_size
if half:
model.half() # to FP16
def attempt_load(weights, map_location=None):
# Loads an ensemble of models weights=[a,b,c] or a single model weights=[a] or weights=a
model = Ensemble()
for w in weights if isinstance(weights, list) else [weights]:
attempt_download(w)
ckpt = torch.load(w, map_location=map_location) # load
model.append(ckpt['ema' if ckpt.get('ema') else 'model'].float().fuse().eval()) # FP32 model
# Compatibility updates
for m in model.modules():
if type(m) in [nn.Hardswish, nn.LeakyReLU, nn.ReLU, nn.ReLU6, nn.SiLU]:
m.inplace = True # pytorch 1.7.0 compatibility
elif type(m) is Conv:
m._non_persistent_buffers_set = set() # pytorch 1.6.0 compatibility
if len(model) == 1:
return model[-1] # return model
else:
print('Ensemble created with %s\n' % weights)
for k in ['names', 'stride']:
setattr(model, k, getattr(model[-1], k))
return model # return ensemble
def check_img_size(img_size, s=32):
# Verify img_size is a multiple of stride s
new_size = make_divisible(img_size, int(s)) # ceil gs-multiple
if new_size != img_size:
print('WARNING: --img-size %g must be multiple of max stride %g, updating to %g' % (img_size, s, new_size))
return new_size
是否进行第二步检测分类，通过resnet101对通过yolov5识别检测后的结果在进行一次预测。
# Second-stage classifier
classify = False
if classify:
modelc = load_classifier(name='resnet101', n=2) # initialize
modelc.load_state_dict(torch.load('weights/resnet101.pt', map_location=device)['model']).to(device).eval()
数据集的载入：
第一行是声明视频保存和vid_writer两个变量。
第二行，if里是如果检测目标是摄像头或者网络视频源，check_imshow()检测运行环境是否是在docker下（docker环境下不能打开界面显示，会报错），然后是视频流或者摄像头的一个载入。
如果不是视频流或者摄像头，载入检测对象为图片、视频或者文件夹内的文件，通过迭代的方式一个一个的进行检测。
# Set Dataloader
vid_path, vid_writer = None, None
if webcam:
view_img = check_imshow()
cudnn.benchmark = True # set True to speed up constant image size inference
dataset = LoadStreams(source, img_size=imgsz, stride=stride)
else:
dataset = LoadImages(source, img_size=imgsz, stride=stride)
载入模型内部的class names，以及当前检测每一个name对应的矩形框的颜色
# 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]
开始执行检测推理部分的代码：
第一个if内的代码是运行一次模型推理，推理为一张全为数值0的图片，图片分辨率为输入的imgsz* imgsz（opencv里的话是全黑的图片）。主要可能是为了检测模型检测是否会报错和让GPU开始预热，哈哈哈。
t0是为了计算整个推理过程的时间的开始计时，即完成文件夹内图像的推理时间。
dataset里的内容可以参考下LoadImages类函数里的__next__()，并通过__iter__()，进行文件夹内部文件的迭代更新。
for里的内容为先将图片内的数值载入到的device（GPU or CPU），如果是在GPU上进行推理（前面的half）将像数值的uint8类型转为fp16或者fp32，然后将像数值的变化范围固定到0.0-1.0之间（对深度学习推理计算更友好吧），最后判断图片的shape是不是3，因为在tensor环境下使用的是ndimension()，如果不是则增加一层。
path为检测对象的地址，img为将检测的原始图片resize到目标尺寸（不失真放缩，通过填充RBG为114, 114, 114的颜色条）
# Run inference
if device.type != 'cpu':
model(torch.zeros(1, 3, imgsz, imgsz).to(device).type_as(next(model.parameters()))) # run once
t0 = time.time()
for path, img, im0s, vid_cap in dataset:
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)
t1是计算单个文件的推理时间的开始计时，在GPU上通过torch.cuda.synchronize()使推理时间计算的更加精确。
pred是模型推理输出的全部结果，以输入尺寸图片为640
640为例，coco80类检测，输出的pred的shape为25200
85，（20
20
3）大分辨率下的20
20个滑块同时乘以大分辨率对应的3个anchor；（40
40
3）中分辨率下的40
40个滑块同时乘以中分辨率对应的3个anchor；（80
80
3）小分辨率下的80*80个滑块同时乘以小分辨率对应的3个anchor；三组数相加为25200，85其中[0:4]为anchor画出来的矩形框，[4]为ojbconf即anchor中含有80类中目标检测的置信度，后面的[5:]为80类别对应的每一类的置信度。
# Inference
t1 = time_synchronized()
pred = model(img, augment=opt.augment)[0]
def time_synchronized():
# pytorch-accurate time
if torch.cuda.is_available():
torch.cuda.synchronize()
return time.time()
极大值抑制，non_max_suppression()函数太长了，就不在这里贴出来了，有兴趣的可以在源码里看看，在这里我大致说下极大值抑制的流程吧。
第一步是判断pred列表里ojbconf也就是pred中其中一个anchor中[4]小于opt.conf_thres的值，返回一个True和Flase的列表，长度为25200。
当遇到判断列表中的第一个True，先计算类别置信度，即obj_conf*80类别算法得出对应的置信度，再将前面4个数即anchor对应的矩形框转换为xyxy的矩形框（box）的像数点坐标值。然后取出得到置信度（conf）中最大的值和所在的位置（classID），再将box的xyxy、conf以及classID拼凑起来为一个列表。最后通过NMS极大值抑制，通过对计算每个为True的box的交并比，输入满足要求的输入阈值的一个list，list的shape[1]为6,[x,y,x,y,conf,classID]。
# Apply NMS
pred = non_max_suppression(pred, opt.conf_thres, opt.iou_thres, classes=opt.classes, agnostic=opt.agnostic_nms)
t2 = time_synchronized()
是否对预测结果进行二次预测分类
# Apply Classifier
if classify:
pred = apply_classifier(pred, modelc, img, im0s)
预测结果的处理：
首先是判断预测的对象是是为视频源还是其他的对象，设声明结果保存的地址（图片结果以及标签结果），s为每张resize到目标尺寸图像的shape，即喂入网络的图片尺寸（str），在tensor上获得原始检测画面的size。
# Process detections
for i, det in enumerate(pred): # detections per image
if webcam: # batch_size >= 1
p, s, im0, frame = path[i], '%g: ' % i, im0s[i].copy(), dataset.count
else:
p, s, im0, frame = path, '', im0s, getattr(dataset, 'frame', 0)
p = Path(p) # to Path
save_path = str(save_dir / p.name) # img.jpg
txt_path = str(save_dir / 'labels' / p.stem) + ('' if dataset.mode == 'image' else f'_{frame}') # img.txt
s += '%gx%g ' % img.shape[2:] # print string
gn = torch.tensor(im0.shape)[[1, 0, 1, 0]] # normalization gain whwh
如果det中有相应的检测目标，将det[:, 0:4]，即在img（原图（im0s）resize到targetSize后的图像）上xyxy对于的像数值坐标还原到原图（im0s）的像数值坐标，为了后续方便在原图上画出目标物的矩形框。
然后是确定det[:, -1]中的所存在同类的classID一共有多少个，并有几个类别，是在后续输出检测完成后打印检测结果的相关信息。
检测结果写入，如果在程序执行时输入了 –save-txt,则将类别 xywh 置信度 写入保存在save_dir / ‘labels’/图片名中，如果去掉置信度和前后的空格就是通过labelimg标注后的yolo标签。
执行程序的时候，没有设置 –nosave或设置了–view-img，则在im0s原始图像上画出预测结果的矩形框，为了后续的保存图像、视频结果或者网络视频源。
在命令框内打印检测总体结果
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):
if save_txt: # Write to file
xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist() # normalized xywh
line = (cls, *xywh, conf) if opt.save_conf else (cls, *xywh) # label format
with open(txt_path + '.txt', 'a') as f:
f.write(('%g ' * len(line)).rstrip() % line + '\n')
if save_img or view_img: # Add bbox to image
label = f'{names[int(cls)]} {conf:.2f}'
plot_one_box(xyxy, im0, label=label, color=colors[int(cls)], line_thickness=3)
# Print time (inference + NMS)
print(f'{s}Done. ({t2 - t1:.3f}s)')
检测时需要查看图片或者视频结果，通过imshow显示，显示检测后的图像或视频结果。
# Stream results
if view_img:
cv2.imshow(str(p), im0)
cv2.waitKey(1) # 1 millisecond
当检测对象是图片是，保存图片到save_path，为save_dir / 完整的图片名（有尾缀）。
当检测对象是视频或者视频源时，因为对视频检测是对视频中的每一帧的图像进行检测，检测视频显示是通过将视频的每一帧拼接成视频显示出来的，当视频检测结果保存的save_path在同一个视频的时候不会发生变化，当发生变化时，cv2.VideoWriter里的save_path也需要相应的改变。
如果检测对象是本地视频则通过cv2获取fps，w，h，检测对象是网络视频默认fps=30。
然后通过cv2.VideoWriter(args).write（im0）写入每一帧检测结果，保存视频。
if save_img:
if dataset.mode == 'image':
cv2.imwrite(save_path, im0)
else: # 'video' or 'stream'
if vid_path != save_path: # new video
vid_path = save_path
if isinstance(vid_writer, cv2.VideoWriter):
vid_writer.release() # release previous video writer
if vid_cap: # video
fps = vid_cap.get(cv2.CAP_PROP_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 = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
vid_writer.write(im0)
打印最总的检测结果。
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 ''
print(f"Results saved to {save_dir}{s}")
print(f'Done. ({time.time() - t0:.3f}s)')
欢迎讨论，并指出纠正上面的错误。