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

zstar

发布于 2022-09-16 09:15:07

3K0

发布于 2022-09-16 09:15:07

文章被收录于专栏：往期博文

背景

在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' ]

注：这里的路径替换成自己的路径即可。