2021SC@SDUSC
目录
请不要忽视代码中的注释
正式推理
for path, img, im0s, vid_cap in dataset:
t1 = time_sync()
# 处理每一张图片/视频的格式
if onnx:
img = img.astype('float32')
else:
img = torch.from_numpy(img).to(device)
img = img.half() if half else img.float() # uint8 to fp16/32
img = img / 255.0 # 0 - 255 to 0.0 - 1.0
if len(img.shape) == 3:
img = img[None] # expand for batch dim
t2 = time_sync()
dt[0] += t2 - t1
nms出去多余的框
# NMS
pred = non_max_suppression(pred, conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det)
dt[2] += time_sync() - t3
# Apply NMS 进行NMS
# conf_thres: 置信度阈值
# iou_thres: iou阈值
# classes: 是否只保留特定的类别 默认为None
# agnostic_nms: 进行nms是否也去除不同类别之间的框
# max_det: 每张图片的最大目标个数 默认1000
# pred: [num_obj, 6] = [5, 6] 这里的预测信息pred还是相对于 img_size(640) 的
考虑进行二次分类
# Second-stage classifier (optional)
# Apply Classifier 如果需要二次分类 就进行二次分类 一般是不需要的
if classify:
pred = apply_classifier(pred, modelc, img, im0s)
后续保存或者打印预测信息
# Process predictions
# 对每张图片进行处理 将pred(相对img_size 640)映射回原图img0 size
for i, det in enumerate(pred): # per image
seen += 1
if webcam: # batch_size >= 1
# 如果输入源是webcam(网页)则batch_size>=1 取出dataset中的一张图片
p, s, im0, frame = path[i], f'{i}: ', im0s[i].copy(), dataset.count
else:
# 但是大部分我们一般都是从LoadImages流读取本都文件中的照片或者视频 所以batch_size=1
# p: 当前图片/视频的绝对路径 如 F:\yolo_v5\yolov5-U\data\images\bus.jpg
# s: 输出信息 初始为 ''
# im0: 原始图片 letterbox + pad 之前的图片
# frame: 初始为0 可能是当前图片属于视频中的第几帧?
p, s, im0, frame = path, '', im0s.copy(), getattr(dataset, 'frame', 0)
# 当前图片路径 如 F:\yolo_v5\yolov5-U\data\images\bus.jpg
p = Path(p) # to Path
# 图片/视频的保存路径save_path 如 runs\\detect\\exp8\\bus.jpg
save_path = str(save_dir / p.name) # img.jpg
# txt文件(保存预测框坐标)保存路径 如 runs\\detect\\exp8\\labels\\bus
txt_path = str(save_dir / 'labels' / p.stem) + ('' if dataset.mode == 'image' else f'_{frame}') # img.txt
# print string 输出信息 图片shape (w, h)
s += '%gx%g ' % img.shape[2:] # print string
# normalization gain gn = [w, h, w, h] 用于后面的归一化
gn = torch.tensor(im0.shape)[[1, 0, 1, 0]] # normalization gain whwh
# imc: for save_crop 在save_crop中使用
imc = im0.copy() if save_crop else im0 # for save_crop
annotator = Annotator(im0, line_width=line_thickness, pil=not ascii)
if len(det):
# Rescale boxes from img_size to im0 size
# 将预测信息(相对img_size 640)映射回原图 img0 size
det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()
# Print results
# 输出信息s + 检测到的各个类别的目标个数
for c in det[:, -1].unique():
# 将每个图片的预测信息分别存入save_dir/labels下的xxx.txt中 每行: class_id+score+xywh
n = (det[:, -1] == c).sum() # detections per class
s += f"{n} {names[int(c)]}{'s' * (n > 1)}, " # add to string
# Write results
# 保存预测信息: txt、img0上画框、crop_img
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 save_conf else (cls, *xywh) # label format
with open(txt_path + '.txt', 'a') as f:
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))
if save_crop:
save_one_box(xyxy, imc, file=save_dir / 'crops' / names[c] / f'{p.stem}.jpg', BGR=True)
# Print time (inference-only)
# 打印前向传播 + NMS 花费的时间
print(f'{s}Done. ({t3 - t2:.3f}s)')
# Stream results
im0 = annotator.result()
# 是否需要显示我们预测后的结果 img0(此时已将pred结果可视化到了img0中)
if view_img:
cv2.imshow(str(p), im0)
cv2.waitKey(1) # 1 millisecond
# Save results (image with detections)
# 是否需要保存图片或视频(检测后的图片/视频 里面已经被我们画好了框的) img0
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)
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)
推理结束,保存结果,打印信息
# Print results
# 保存预测的label信息 xywh等 save_txt
t = tuple(x / seen * 1E3 for x in dt) # speeds per image
print(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 ''
print(f"Results saved to {colorstr('bold', save_dir)}{s}")
if update:
# strip_optimizer函数将optimizer从ckpt中删除 更新模型
strip_optimizer(weights) # update model (to fix SourceChangeWarning)