推理阶段是整个检测模型完成后,要对模型进行测试的部分。很重要的一部分,只有了解了这个部分,才能在比赛或者项目提交中很好的输出自己模型的检测结果。同时,推理输出对模型部署在不同的环境下也是十分重要的。
源码:https://github.com/ultralytics/yolov5
版本yolov5 v6.1
@torch.no_grad() # 装饰器,推理部分不进行反向传播计算
def run(
weights=ROOT / 'yolov5s.pt', # 加载的训练模型权重路径
source=ROOT / 'data/images', # 输入图像或视频文件路径
data=ROOT / 'data/coco128.yaml', # dataset.yaml path,自己生成的数据集yaml文件路径
imgsz=(640, 640), # inference size (height, width)
conf_thres=0.25, # confidence threshold
iou_thres=0.45, # NMS IOU threshold
max_det=1000, # maximum detections per image,默认一张图最多输出1000个目标
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,保存成.txt文件
save_conf=False, # save confidences in --save-txt labels,是否将置信度保存到txt文件
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,进行nms时是否将其他类当作同一个类别处理
augment=False, # augmented inference,是否进行测试推理时的数据增强,TTA之类的
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
):
source = str(source) # 图像输入路径
save_img = not nosave and not source.endswith('.txt') # save inference images
is_file = Path(source).suffix[1:] in (IMG_FORMATS + VID_FORMATS) # 判断文件的后缀名是否是图像或视频
is_url = source.lower().startswith(('rtsp://', 'rtmp://', 'http://', 'https://')) # 是否是网址
webcam = source.isnumeric() or source.endswith('.txt') or (is_url and not is_file) #视频流文件
if is_url and is_file:
source = check_file(source) # download
# Directories
save_dir = increment_path(Path(project) / name, exist_ok=exist_ok) # increment run,默认路径runs/detect/exp
(save_dir / 'labels' if save_txt else save_dir).mkdir(parents=True, exist_ok=True) # make dir
# Load model
device = select_device(device) # 选择设备
# 不同的权重格式选择不同的加载方式 ,如.pt,.onnx,.trt
model = DetectMultiBackend(weights, device=device, dnn=dnn, data=data, fp16=half)
stride, names, pt = model.stride, model.names, model.pt
imgsz = check_img_size(imgsz, s=stride) # check image size,检查选择输入模型的尺寸是否符合32的倍数
整个数据处理的步骤;
# Dataloader
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, auto=pt)
bs = len(dataset) # batch_size
else:
# 这里只是简单的加载输入图像数据的路径
dataset = LoadImages(source, img_size=imgsz, stride=stride, auto=pt)
bs = 1 # batch_size
vid_path, vid_writer = [None] * bs, [None] * bs
# Run inference
# 先进行预热,用全为0的先进行一次推理,过一遍代码
model.warmup(imgsz=(1 if pt else bs, 3, *imgsz)) # warmup
# 这里迭代的时候要进入,utils/dataloaders.py里的__next__里进行letterbox数据宽高放缩处理
for path, im, im0s, vid_cap, s in dataset:
t1 = time_sync() # 计算时间
im = torch.from_numpy(im).to(device) # 转成tensor并将数据转到cuda上
im = im.half() if model.fp16 else im.float() # uint8 to fp16/32,选择使用的精度,半浮点型速度要更快
im /= 255 # 0 - 255 to 0.0 - 1.0 这是像素值归一化,跟坐标归一化无关
if len(im.shape) == 3:
im = im[None] # expand for batch dim,拓展一个batch维度
t2 = time_sync()
dt[0] += t2 - t1 # 计算数据预处理时间
# Padded resize
img = letterbox(img0, self.img_size, stride=self.stride, auto=self.auto)[0]
---------------------------------------------------------------------------
---------------------------------------------------------------------------
def letterbox(im, new_shape=(640, 640), color=(114, 114, 114), auto=True, scaleFill=False, scaleup=True, stride=32):
# Resize and pad image while meeting stride-multiple constraints
shape = im.shape[:2] # 输入原图高宽[height, width]
if isinstance(new_shape, int):
new_shape = (new_shape, new_shape)
# Scale ratio (new / old),求自己要求输入模型图像的高宽和原图高宽的最小比值
r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])
if not scaleup: # only scale down, do not scale up (for better val mAP)
r = min(r, 1.0) # 这表示如果比值大于1不进行比例缩放
# Compute padding
ratio = r, r # width, height ratios
# 对原图宽高进行最短边比列缩放,保持宽高比不变,使图像尽量不失真
new_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))
# 为了符合32的倍数要求(满足最小预测输出的特征层),要对缩放的尺寸进行padding
dw, dh = new_shape[1] - new_unpad[0], new_shape[0] - new_unpad[1] # wh padding
if auto: # minimum rectangle
dw, dh = np.mod(dw, stride), np.mod(dh, stride) # wh padding np.mod取模运算
elif scaleFill: # stretch
dw, dh = 0.0, 0.0
new_unpad = (new_shape[1], new_shape[0])
ratio = new_shape[1] / shape[1], new_shape[0] / shape[0] # width, height ratios
dw /= 2 # divide padding into 2 sides
dh /= 2 # padding是对图像两边进行的
if shape[::-1] != new_unpad: # 和原图宽高不等,opencv进行resize
im = cv2.resize(im, new_unpad, interpolation=cv2.INTER_LINEAR)
top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1)) # 四舍五入后取整
left, right = int(round(dw - 0.1)), int(round(dw + 0.1))
# padding的部分是无效区域,用color=(114, 114, 114)常数值填充,前面讲mmdetection里的mask记录放缩图像无效区域一个意思
im = cv2.copyMakeBorder(im, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color) # add border
# 返回缩放的宽高,比例和padding大小,方便后面对归一化的预测输出进行还原,推理只要im
return im, ratio, (dw, dh)
---------------------------------------------------------------------------------
# opncv读取的图片数据是HWC和BGR格式的,为了方便要进行处理
# Convert
img = img.transpose((2, 0, 1))[::-1] # HWC to CHW, BGR to RGB
img = np.ascontiguousarray(img)
return path, img, img0, self.cap, s
推理步骤:
# Inference
visualize = increment_path(save_dir / Path(path).stem, mkdir=True) if visualize else False
# [b,num_pre_boxes,(5+classes)]这里是输出归一化的boxes:x,y,h,w,5表示[x,y,h,w,conf]
pred = model(im, augment=augment, visualize=visualize)
t3 = time_sync()
dt[1] += t3 - t2 # 模型预测输出时间
# 多类别的nms,部署的话要可以用opencv的,也可以用c++可以自己写
# NMS,经过下面的NMS后才是[num_P,6],[x1,y1,x2,y2,conf,cls]
pred = non_max_suppression(pred, conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det)# [x1,y1,x2,y2,conf,cls]
dt[2] += time_sync() - t3 # NMS的时间
# Process predictions
for i, det in enumerate(pred): # per image
seen += 1
if webcam: # batch_size >= 1
p, im0, frame = path[i], im0s[i].copy(), dataset.count
s += f'{i}: '
else:
p, im0, frame = path, im0s.copy(), getattr(dataset, 'frame', 0)
p = Path(p) # to Path
save_path = str(save_dir / p.name) # im.jpg
txt_path = str(save_dir / 'labels' / p.stem) + ('' if dataset.mode == 'image' else f'_{frame}') # im.txt
s += '%gx%g ' % im.shape[2:] # print string
gn = torch.tensor(im0.shape)[[1, 0, 1, 0]] # normalization gain whwh, 获取原图大小
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,im:640,640,im0:原图大小
det[:, :4] = scale_coords(im.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
# xyxy2xywh:字面意思,将左上右下角的坐标值,转换为中心宽高值,再除以比例,归一化保存到txt文件上
xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist() # normalized xywh
# txt文件里每行记录(cls, *xywh),如果要记录置信度,将save_conf设置为true
line = (cls, *xywh, conf) if save_conf else (cls, *xywh) # label format
# 打开txt文件,将推理的结果写入
with open(f'{txt_path}.txt', 'a') as f:
f.write(('%g ' * len(line)).rstrip() % line + '\n')
# 后面的代码都是保存和展示结果文件的了,这里就不解析了
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)
在前面的博客,讲了单类别的nms,但是如果是多类别的nms可以怎么做呢?
已知的方法:
yolov5就是采用的第三种方法,这里这要讲下主要的代码,完整代码在utils/general.py里的non_max_suppression里
# 主要代码
max_wh = 7680 # 最大宽高
----------------------------------------------------------
# 这里的代码主要解释x的维度代表什么,下面有用到
# 找出的conf_thres(类别分数*置信度)最大的对应的一个类和索引,j是类别label的索引:[0,1,1,2]这种
conf, j = x[:, 5:].max(1, keepdim=True)
# 在维度1上重新合并[经过初步分数阈值筛选的预测数量,4+1+1=6],取出大于conf_thres阈值的预测值[k,6]
x = torch.cat((box, conf, j.float()), 1)[conf.view(-1) > conf_thres]
---------------------------------------------------------
# 进行多类别的nms
# Batched NMS,如果agnostic为true,表示所有的作为一类进行nms,默认false,每类乘上一个最大wh长度的偏移值
c = x[:, 5:6] * (0 if agnostic else max_wh) # classes
# 之后在boxes上加上偏移量作为iou计算的boxes
boxes, scores = x[:, :4] + c, x[:, 4] # boxes (offset by class), scores
i = torchvision.ops.nms(boxes, scores, iou_thres) # NMS
推理输出部分的理解是十分重要的,在比赛和项目中有提交结果格式的要求,都要在这修改相应的代码。同时,也可以在推理过程中使用一些trick去提升性能。同时,为了实时性,算法的开发部署主要就是推理阶段的部署。为了加速推理速度,提升检测性能,可以很好的将训练好的torch模型转换成onnx或者tensorrt,使用c++进行部署。
yolov5使用自己的数据集相关代码
yolov5网络结构代码解读
yolov5的正负样本的定义和匹配