【目标检测】YOLOv5跑通VisDrone数据集

背景

在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！