yolov5的推理输出detect.py部分-Toy模板网

这篇具有很好参考价值的文章主要介绍了yolov5的推理输出detect.py部分。希望对大家有所帮助。如果存在错误或未考虑完全的地方，请大家不吝赐教，您也可以点击"举报违法"按钮提交疑问。

前言

推理阶段是整个检测模型完成后，要对模型进行测试的部分。很重要的一部分，只有了解了这个部分，才能在比赛或者项目提交中很好的输出自己模型的检测结果。同时，推理输出对模型部署在不同的环境下也是十分重要的。
源码：https://github.com/ultralytics/yolov5
版本yolov5 v6.1

detect.py

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
):

2.设置文件保存路径

	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

3.加载训练好的模型权重

# 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的倍数

4.数据处理

整个数据处理的步骤;

1.先对输入图像的原图进行最短边的放缩;
2.为了匹配yolov5的下采样操作还要使得输入模型的宽高符合最后一次下采样stride(三层一般32，四层64)的倍数，还要对其进行padding。
+3. 同时，因为opencv读取图像是BGR格式，要将其转化为RGB格式。
4.最后，还要对输入的像素值进行归一化。

# 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  # 计算数据预处理时间

letterbox进行resize和pad，opencv读取的bgr转成rgb

# 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

5.进行推理

推理步骤：

1.将输入图像经过模型得到预测输出；
2.对预测输出进行NMS处理；
3.对NMS处理后的Boxes进行相应的后处理。

     # Inference
     visualize = increment_path(save_dir / Path(path).stem, mkdir=True) if visualize else False
     # [b,num_pre_boxes,(5+cls_scores)]这里是输出归一化的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)

6.yolov5里的nms

在前面的博客，讲了单类别的nms，但是如果是多类别的nms可以怎么做呢？
已知的方法：

第一种方法就是，对循环遍历每类，然后每类进行nms；
第二种方法，直接对所有的box进行分数降序排列，之后在循环里进行判断，只要类别相同的才进行iou计算，然后再比较阈值。这要用一个标记数组来记录哪些是要筛选的，之后筛选掉。
第三种方法，使用一个偏移量，将不同的类加上一个偏移量，将每个类的所有box单独变换到不同的坐标域，直接按原来的进行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或者engine，使用c++进行部署。