背景
在YOLOv5官方的6.1版本中，我发现Data目录下多了5种数据集，包括无人机航拍数据集VisDrone和遥感数据集xView，这反映了官方也是在小目标检测中在偷偷发力。 这篇博文就来记录如何用YOLOv5跑通VisDrone数据集。
数据集我已上传到网盘里，有需要的读者可以进行下载： https://pan.baidu.com/s/1UNQlZGHZlAZs412tbnpAxg?pwd=8888
数据集结构
数据集总共分四个文件：
- VisDrone2019-DET-train：训练集
- VisDrone2019-DET-val：验证集
- VisDrone2019-DET-test-dev：测试集(带标签)
- VisDrone2019-DET-test-challenge：测试集(不带标签）
不带标签的测试集是用于VisDrone2021挑战赛的，不参加竞赛就暂时不需要用到。
数据集处理
和VOC数据集一样，VisDrone的数据标注是xml形式，需要把它转换成YOLOv5所需的txt格式
在官方提供的VisDrone.yaml下方，已经提供了数据处理的脚本，以它为基础进行简单修改。
在根目录下创建visdrone2yolo.py：
from utils.general import download, os, Path
def visdrone2yolo(dir):
from PIL import Image
from tqdm import tqdm
def convert_box(size, box):
# Convert VisDrone box to YOLO xywh box
dw = 1. / size[0]
dh = 1. / size[1]
return (box[0] box[2] / 2) * dw, (box[1] box[3] / 2) * dh, box[2] * dw, box[3] * dh
(dir / 'labels').mkdir(parents=True, exist_ok=True) # make labels directory
pbar = tqdm((dir / 'annotations').glob('*.txt'), desc=f'Converting {dir}')
for f in pbar:
img_size = Image.open((dir / 'images' / f.name).with_suffix('.jpg')).size
lines = []
with open(f, 'r') as file: # read annotation.txt
for row in [x.split(',') for x in file.read().strip().splitlines()]:
if row[4] == '0': # VisDrone 'ignored regions' class 0
continue
cls = int(row[5]) - 1 # 类别号-1
box = convert_box(img_size, tuple(map(int, row[:4])))
lines.append(f"{cls} {' '.join(f'{x:.6f}' for x in box)}n")
with open(str(f).replace(os.sep 'annotations' os.sep, os.sep 'labels' os.sep), 'w') as fl:
fl.writelines(lines) # write label.txt
dir = Path('D:/Desktop/Work/Dataset/VisDrone') # datasets文件夹下Visdrone2019文件夹目录
# Convert
for d in 'VisDrone2019-DET-train', 'VisDrone2019-DET-val', 'VisDrone2019-DET-test-dev':
visdrone2yolo(dir / d) # convert VisDrone annotations to YOLO labels运行完之后，可以看到在每个数据集下生成了和images对应的labels
创建自己的数据路径文件
在data下面创建mydata.yaml(也可以直接修改VisDrone.yaml)
输入下面内容：
train: D:/Dataset/VisDrone/VisDrone2019-DET-train/images # train images (relative to 'path') 6471 images
val: D:/Dataset/VisDrone/VisDrone2019-DET-val/images # val images (relative to 'path') 548 images
test: D:/Dataset/VisDrone/VisDrone2019-DET-test-dev/images # test images (optional) 1610 images
# Classes
nc: 10 # number of classes
names: [ 'pedestrian', 'people', 'bicycle', 'car', 'van', 'truck', 'tricycle', 'awning-tricycle', 'bus', 'motor' ]注：这里的路径替换成自己的路径即可。
开始训练
在train.py中，主要需要修改下方红框框出来的几个量。 batch-size和workers两个参数根据自己的电脑配置灵活选取。 值得注意的是YOLOv5 5.0版本的输入图片默认尺寸大小是640，在YOLOv5 6.1版本中，默认尺寸修改成了960，这也反映了官方应对小目标检测的策略。
设置完之后，开始训练即可。
视频检测
模型训练完之后，运行detect.py可以对图片或视频进行检测。
使用时，主要修改前三个参数，意义和train.py中类似。
另外，对于密集的小目标，输出的时候默认生成数据标签和置信度可能会对目标产生遮挡，影响观感。
可以修改下方hide-labels和hide-conf两个属性，这样就可以把标签和置信度进行隐藏。
Enjoy yourself！
