Press "Enter" to skip to content

YOLOv5 小目标检测

1.数据集处理

 

(1)数据集下载

 

使用的是VisDrone2021数据集,这是无人机视觉目标检测的数据集,VisDrone2021版本和VisDrone2019是同一个数据集,里面包含很多的小目标。这个数据集不需要注册,可以免费下载。作者提供了百度网盘和谷歌网盘的下载方式。

 

下载地址: Object Detection – VisDrone

 

 

(2)密集区域过滤

 

整理数据集的时候发现有一个分类是 “ignored regions”,这个是忽略的区域,因为有些区域包含了密集的很小的目标,是无法进行标注的,所以我们要把这个区域忽视掉。这里我直接用opencv来进行覆盖住就可以了。

 

 

 

 

(3)图片分割、标签生成

 

因为我们要检测的图像分辨率很大,比如无人机拍摄的图片尺寸为5630×4314,但又有些目标很小,如果直接把图像缩放到640×640训练的话,效果不好,很多小目标就检测不到了。

 

因为:yolov5使用了5次下采样,最后输出的特征图大小是 20 * 20, 40 * 40, 80 * 80 。

 

80* 80 负责检测小目标的,对应到640 * 640 上,每格特征图对应的感受野大小是 640/80=8 * 8。再对应到原图中,以长边为例,5630/640 * 8 = 71,即原图中目标小于71像素的目标,是无法学习到有效特征的。

 

所以要对原图分割成多个小图再进行检测,我这里把图像分割成 2行3列,即6个小图。

 

有一点值得注意的是,有一些目标正好位于两个小图中间,正好被截断了,这就可能导致检测不到目标。为了避免这种情况,我们 在两个小图之间设置一个overlap重叠区域 ,这里我设置的重叠区域面积占总面积的20%。

 

然后把标签中boxes的坐标也要跟着做相应的变换,保存到 VisDrone_chip 文件夹中。同时,生成yolov5训练所需要的标签格式。

 

原图分割后生成的训练目录:

 

 

2.模型训练

 

(1)创建自己的数据

 

我们首先用git拉取YOLOv5的代码,github地址:

 

GitHub – ultralytics/yolov5: YOLOv5  in PyTorch > ONNX > CoreML > TFLite YOLOv5  in PyTorch > ONNX > CoreML > TFLite. Contribute to ultralytics/yolov5 development by creating an account on GitHub. https://github.com/ultralytics/yolov5.git

 

然后在 yolov5/data/ 目录下创建自己的数据集配置文件 VisDrone_data.yaml :

 

# VisDrone2019-DET dataset https://github.com/VisDrone/VisDrone-Dataset
# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
train: ./data/VisDrone_chip/images/train  # train images (relative to 'path')  6471 images
val: ./data/VisDrone_chip/images/val  # val images (relative to 'path')  548 images
test: ./data/VisDrone_chip/images/test  # test images (optional)  1610 images
# Classes
nc: 10  # number of classes
names: ['pedestrian', 'people', 'bicycle', 'car', 'van', 'truck', 'tricycle', 'awning-tricycle', 'bus', 'motor']

 

这里我们把 “others” 这个分类去掉,一共10个类别。

 

(2)创建自己的模型

 

进入到 yolov5/models/ 目录下,拷贝一个模型作为自己的模型,这里我分别尝试了YOLOv5的small版本和large版本模型,这里就用small版本进行讲解。

 

cp yolov5s.yaml yolov5s_visdrone.yaml

 

然后修改类别的数量,这里共有10个分类,把nc的参数改为10:

 

 

(3)模型训练

 

下载pretrain模型模型 YOLOv5s , 存放在yolov5 目录里。

 

开始训练:

 

python train.py --img 640 --batch 16 --epochs 100 --data ./data/VisDrone_data.yaml --weights yolov5s.pt

 

这里我训练100个epochs,经过9个多小时的训练,训练结果如下:

 

 

 

[email protected] 达到了0.591,效果挺好的。

 

YOLOv5l 版本则经过漫长的26个小时训练, [email protected] 达到了0.648,效果非常好。

 

3.推理合并

 

在模型推理的时候,输入的是一张无人机拍下是原始图像,我们同样也需要把原图切割成多个小图来推理,再把小图的推理的结果合并到原图,然后再统一做nms操作。

 

步骤:

 

(1)小图跑模型推理,得到推理结果pred;

 

(2)对pred的结果中boxes位置进行坐标转换,转换为对应于原图中的位置;

 

(3)把各个小图推理的结果用 torch.cat 来进行合并;

 

(4)使用nms非极大值抑制 过滤掉重复的框。

 

测试效果:

 

 

 

 

 

可以看到效果是非常好的。

Be First to Comment

发表回复

您的电子邮箱地址不会被公开。 必填项已用*标注