yolov5训练VisDrone数据集

最新推荐文章于 2024-10-09 12:09:53 发布

wyw0000

最新推荐文章于 2024-10-09 12:09:53 发布

阅读量607

点赞数 5

分类专栏： yolo 文章标签： YOLO

本文链接：https://blog.csdn.net/wyw0000/article/details/136895297

版权

yolo 专栏收录该内容

11 篇文章 0 订阅

订阅专栏

本文介绍了如何从VisDrone数据集下载和转换标签，修改YOLov5配置文件以适应该数据集，进行模型训练，以及如何将模型转换为TensorRT引擎以支持目标推理。

摘要由CSDN通过智能技术生成

1. 数据集、模型、代码下载

下载地址
 yolov5下载
 yolov5-tensorRT下载
 yolov5模型下载

2. 数据集标签转换

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


if __name__ == '__main__':
    dir = Path('E:\YOLO-datasets\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

3. 修改配置文件

3.1 修改数据文件

配置文件位于yolov5\data\VisDrone.yaml

train: ../datasets/images/train  # train images (relative to 'path')  6471 images
val: ../datasets/images/val  # val images (relative to 'path')  548 images
test: ../datasets/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']

train、val、test修改为VisDrone数据集下载解压后的目录

3.2 修改模型文件

比如：使用yolov5l模型，那么修改yolov5\models目录下的yolov5l.yaml
拷贝yolov5l.yaml并明明为yolov5l-VisDrone.yaml

nc: 10  # number of classes
depth_multiple: 1.0  # model depth multiple
width_multiple: 1.0  # layer channel multiple

修改nc为10，即VisDrone数据集的分类个数

4. 训练

下载预训练权重，开始训练

python train.py --weights yolov5l.pt --data data/VisDrone.yaml --img 640--device 0 --cfg models/yolov5l-VisDrone.yaml --batch-size 16 --epochs 100

5. 模型转换

如果要使用tensorRT进行目标推理，则需要将pytorch训练的模型pt格式转换为engine。

5.1 转换为wts

python gen_wts.py -w yolov5l.pt -o yolov5l.wts

5.2 编译yolov5

这里不详述
注意：
CLASS_NUM 、INPUT_H 、INPUT_W 的值一定要和模型对应上，比如：VisDrone是10分类，yolov5l模型输入时640，那么CLASS_NUM = 10，INPUT_H = 640；INPUT_W = 640；

    static constexpr int CLASS_NUM = 10;
    static constexpr int INPUT_H = 1280;  // yolov5's input height and width must be divisible by 32.
    static constexpr int INPUT_W = 1280;