目标检测YOLOv5跑通VisDrone数据集

Posted zstar-_

tags:

篇首语:本文由小常识网(cha138.com)小编为大家整理,主要介绍了目标检测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-labelshide-conf两个属性,这样就可以把标签和置信度进行隐藏。

Enjoy yourself!

以上是关于目标检测YOLOv5跑通VisDrone数据集的主要内容,如果未能解决你的问题,请参考以下文章

目标检测YOLOv5跑通VOC2007数据集(修复版)

目标检测YOLOv6理论解读+实践测试VisDrone数据集

目标检测TPH-YOLOv5:基于transformer的改进yolov5的无人机目标检测

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

目标检测YOLOv5:640与1280分辨率效果对比

目标检测YOLOR理论简介+实践测试VisDrone数据集