【目标检测】YOLOv5针对小目标检测的改进模型/添加帧率检测

问题背景

众所周知，YOLOv5会对输入的图片进行放缩，并进行32倍下采样。对于一些分辨率很高的遥感/无人机图片，小目标难以被训练识别。
本篇博文就来尝试这篇博文YOLOV5 模型和代码修改——针对小目标识别所提到的一种改进方案。

我所使用的是YOLOv5-5.0版本，数据集采用VisDrone数据集。

检测头改进

模型方面的修改：作者在模型上增加了一个更小的Anchor并添加了一个更小的检测头。
yolov5l_modify.yaml

# parameters
nc: 10  # number of classes
depth_multiple: 1.0  # model depth multiple
width_multiple: 1.0  # layer channel multiple

# anchors
anchors:
  - [5,6, 8,14, 15,11]  #4
  - [10,13, 16,30, 33,23]  # P3/8
  - [30,61, 62,45, 59,119]  # P4/16
  - [116,90, 156,198, 373,326]  # P5/32

# YOLOv5 backbone
backbone:
  # [from, number, module, args]
  [[-1, 1, Focus, [64, 3]],  # 0-P1/2
   [-1, 1, Conv, [128, 3, 2]],  # 1-P2/4
   [-1, 3, BottleneckCSP, [128]],   #160*160
   [-1, 1, Conv, [256, 3, 2]],  # 3-P3/8
   [-1, 9, BottleneckCSP, [256]],  #80*80
   [-1, 1, Conv, [512, 3, 2]],  # 5-P4/16
   [-1, 9, BottleneckCSP, [512]], #40*40
   [-1, 1, Conv, [1024, 3, 2]],  # 7-P5/32
   [-1, 1, SPP, [1024, [5, 9, 13]]],
   [-1, 3, BottleneckCSP, [1024, False]],  # 9   20*20
  ]

# YOLOv5 head
head:
  [[-1, 1, Conv, [512, 1, 1]],  #20*20
   [-1, 1, nn.Upsample, [None, 2, 'nearest']], #40*40
   [[-1, 6], 1, Concat, [1]],  # cat backbone P4  40*40
   [-1, 3, BottleneckCSP, [512, False]],  # 13     40*40

   [-1, 1, Conv, [512, 1, 1]], #40*40
   [-1, 1, nn.Upsample, [None, 2, 'nearest']],
   [[-1, 4], 1, Concat, [1]],  # cat backbone P3   80*80
   [-1, 3, BottleneckCSP, [512, False]],  # 17 (P3/8-small)  80*80

   [-1, 1, Conv, [256, 1, 1]], #18  80*80
   [-1, 1, nn.Upsample, [None, 2, 'nearest']], #19  160*160
   [[-1, 2], 1, Concat, [1]], #20 cat backbone p2  160*160
   [-1, 3, BottleneckCSP, [256, False]], #21 160*160

   [-1, 1, Conv, [256, 3, 2]],  #22   80*80
   [[-1, 18], 1, Concat, [1]], #23 80*80
   [-1, 3, BottleneckCSP, [256, False]], #24 80*80

   [-1, 1, Conv, [256, 3, 2]], #25  40*40
   [[-1, 14], 1, Concat, [1]],  # 26  cat head P4  40*40
   [-1, 3, BottleneckCSP, [512, False]],  # 27 (P4/16-medium) 40*40

   [-1, 1, Conv, [512, 3, 2]],  #28  20*20
   [[-1, 10], 1, Concat, [1]],  #29 cat head P5  #20*20
   [-1, 3, BottleneckCSP, [1024, False]],  # 30 (P5/32-large)  20*20

   [[21, 24, 27, 30], 1, Detect, [nc, anchors]],  # Detect(p2, P3, P4, P5)
  ]

模型方面的改进有点类似于TPH-YOLOv5。

图像切割

作者在检测的时候(detect.py)增加了一个图像切分的步骤，即将大图切分成各个小块，分别进行检测，然后再进行融合。
增加的代码如下：

 # Inference
 t1 = time_sync()
 # pred = model(img, augment=opt.augment)[0] 原始

 '''
 此处进行改进
 '''
 mulpicplus = "3"  # 1 for normal,2 for 4pic plus,3 for 9pic plus and so on
 assert (int(mulpicplus) >= 1)
 if mulpicplus == "1":
     pred = model(img,
                  augment=augment,
                  visualize=increment_path(save_dir / Path(path).stem, mkdir=True) if visualize else False)[0]
 else:
     xsz = img.shape[2]
     ysz = img.shape[3]
     mulpicplus = int(mulpicplus)
     x_smalloccur = int(xsz / mulpicplus * 1.2)
     y_smalloccur = int(ysz / mulpicplus * 1.2)
     for i in range(mulpicplus):
         x_startpoint = int(i * (xsz / mulpicplus))
         for j in range(mulpicplus):
             y_startpoint = int(j * (ysz / mulpicplus))
             x_real = min(x_startpoint + x_smalloccur, xsz)
             y_real = min(y_startpoint + y_smalloccur, ysz)
             if (x_real - x_startpoint) % 64 != 0:
                 x_real = x_real - (x_real - x_startpoint) % 64
             if (y_real - y_startpoint) % 64 != 0:
                 y_real = y_real - (y_real - y_startpoint) % 64
             dicsrc = img[:, :, x_startpoint:x_real,
                      y_startpoint:y_real]
             '''
             可选，查看切片图片内容
             img2 = dicsrc.squeeze(0).cpu().numpy()
             img2 = img2.transpose((1, 2, 0))
             cv2.imshow('123', img2)
             cv2.waitKey(1)  
             '''
             pred_temp = model(dicsrc,
                               augment=augment,
                               visualize=increment_path(save_dir / Path(path).stem,
                                                        mkdir=True) if visualize else False)[0]
             pred_temp[..., 0] = pred_temp[..., 0] + y_startpoint
             pred_temp[..., 1] = pred_temp[..., 1] + x_startpoint
             if i == 0 and j == 0:
                 pred = pred_temp
             else:
                 pred = torch.cat([pred, pred_temp], dim=1)

代码中注释的部分可将切割的图片可视化展示。

效果检测

为了检测这样做是否有效，我使用改进前的YOLOv5l模型和改进后的YOLOv5l模型对VisDrone数据集训练100个epoch，并挑选了VisDrone测试集中的两张角度较高的图片进行检测，结果如下：

左侧是改进前，右侧是改进后：


通过对比发现两者实际上并没有太大的差异，可能是由于VisDrone数据集拍摄高度还是比较低，无法显示出效果，有待尝试更高分辨率的图片。

帧率检测

在尝试视频检测时，我想到如果能在输出视频中显示帧率就好了。
要实现这个功能只需要在detect.py中插入

# 函数开头插入
tt = time.time()
....

# 添加帧率检测
cv2.putText(im0, "FPS:{:.1f}".format(1. / (time.time() - tt)), (50, 50), cv2.FONT_HERSHEY_SIMPLEX, 2, (0, 0, 235), 4)
tt = time.time()

# 插入在这个位置
if save_img：
	...

代码备份

改进后的代码：https://www.aliyundrive.com/s/SyPtH4N2jcD