【超详细】基于YOLOv8训练无人机视角Visdrone2019数据集

主要内容如下：

1、Visdrone2019数据集介绍
2、下载、制作YOLO格式训练集
3、模型训练及预测
4、Onnxruntime推理

运行环境：Python=3.8（要求>=3.8），torch1.12.0+cu113（要求>=1.8），onnxruntime-gpu==1.12.0
原始数据集百度AI stduio下载链接：https://aistudio.baidu.com/datasetdetail/115729
Visdrone-YOLO格式数据集下载链接：https://aistudio.baidu.com/datasetdetail/295374
训练资源占用：bacth=16，workers=8，yolov8s显存需16G，bacth=8的话8G够用，RTX4080大约1min一个epoch。

往期内容：

【超详细】跑通YOLOv8之深度学习环境配置1-Anaconda安装
【超详细】跑通YOLOv8之深度学习环境配置2-CUDA安装
【超详细】跑通YOLOv8之深度学习环境配置3-YOLOv8安装
【超详细】基于YOLOv8的PCB缺陷检测
【超详细】基于YOLOv8改进1-Drone-YOLO复现

1 数据集介绍

1.1 简介

VisDrone数据集是由天津大学等团队开源的一个大型无人机视角的数据集，官方提供的数据中训练集是6471、验证集是548、测试集1610张。数据集共提供了以下12个类，分别是：‘忽略区域’, ‘pedestrian’, ‘people’, ‘bicycle’, ‘car’, ‘van’,‘truck’, ‘tricycle’, ‘awning-tricycle’, ‘bus’, ‘motor’, ‘others’，其中忽略区域、others是非有效目标区域，本项目中予以忽略；

1.2 示例

1.3 标签格式

**标签含义：**
1. 边界框左上角的x坐标
2. 边界框左上角的y坐标
3. 边界框的宽度
4. 边界框的高度
5. GROUNDTRUTH文件中的分数设置为1或0。1表示在计算中考虑边界框，而0表示将忽略边界框。
6.  类别：忽略区域（0）、行人（1）、人（2）、自行车（3）、汽车（4）、面包车（5）、卡车（6）、三轮车（7）、雨篷三轮车（8）、公共汽车（9）、摩托车（10），其他（11）。
7. GROUNDTRUTH文件中的得分表示对象部分出现在帧外的程度（即，无截断=0（截断比率0%），部分截断=1（截断比率1%°´50%））。
8. GROUNDTRUTH文件中的分数表示被遮挡的对象的分数（即，无遮挡=0（遮挡比率0%），部分遮挡=1（遮挡比率1%°´50%），重度遮挡=2（遮挡率50%~100%））。

2 下载和制作YOLO格式数据集

2.1 下载原始数据集

百度AI stduio下载链接：https://aistudio.baidu.com/datasetdetail/115729
注意：可直接下载已完成转换的YOLO格式数据进行训练，可跳过该阶段，直接训练！链接为：https://aistudio.baidu.com/datasetdetail/295374

下载解压：

注意：由于格式不是YOLO直接可以训练的格式，所以需进行转换！！！

2.2 制作YOLO格式数据集

（1）新建visdrone2yolo.py脚本，脚本内容如下：
（2）修改路径参数–dir_path的值，即自己下载路径；
（2）结果会在原始每个文件夹下生成一个label文件夹，即YOLO格式标签；

import os
from pathlib import Path
import argparsedef visdrone2yolo(dir):from PIL import Imagefrom tqdm import tqdmdef convert_box(size, box):# Convert VisDrone box to YOLO xywh boxdw = 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 directorypbar = tqdm((dir / 'annotations').glob('*.txt'), desc=f'Converting {dir}')for f in pbar:img_size = Image.open((dir / 'images' / f.name).with_suffix('.jpg')).sizelines = []with open(f, 'r') as file:  # read annotation.txtfor row in [x.split(',') for x in file.read().strip().splitlines()]:if row[4] == '0':  # VisDrone 'ignored regions' class 0continuecls = int(row[5]) - 1  # 类别号-1box = 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.txtif __name__ == '__main__':# Create an argument parser to handle command-line argumentsparser = argparse.ArgumentParser()parser.add_argument('--dir_path', type=str, default=r'E:\datasets\visdrone2019', help='visdrone数据集路径')args = parser.parse_args()dir = Path(args.dir_path)# Convertfor d in 'VisDrone2019-DET-train', 'VisDrone2019-DET-val', 'VisDrone2019-DET-test-dev':visdrone2yolo(dir / d)  # convert VisDrone annotations to YOLO labels

3 模型训练及预测

3.1 模型训练

3.1.1 修改数据集配置文件

文件路径：ultralytics-main\ultralytics\cfg\datasets\VisDrone.yaml

3.1.2 创建模型训练脚本

（1）训练方式1-脚本训练
在ultralytics-main目录新建一个train.py脚本，内容如下：
注意：如爆显存，降低batch大小！！！
【如下配置显存需16G，bacth=8的话8G够用，RTX4080大约1min一个epoch】

from ultralytics import YOLOif __name__ == '__main__':# Load a model# model = YOLO("yolov8n.yaml")  # build a new model from scratchmodel = YOLO("yolov8s.pt")  # load a pretrained model (recommended for training)# Use the modelmodel.train(data="VisDrone.yaml", imgsz=640, batch=16, workers=8, cache=True, epochs=100)  # train the modelmetrics = model.val()  # evaluate model performance on the validation set# results = model("ultralytics\\assets\\bus.jpg")  # predict on an imagepath = model.export(format="onnx", opset=13)  # export the model to ONNX format

（2）训练方式2-终端命令行

cd ../ultralytics-main
yolo task=detect mode=train model=yolov8s.pt data=ultralytics/cfg/datasets/VisDrone.yaml batch=16 epochs=100 imgsz=640 workers=8 cache=True device=0

3.1.3 数据分布情况可视化

特点：类别不均衡、小目标较多（640*640输入精度不会太高，可提高输入分辨率，如1280、1536等）。

3.1.4 训练结果可视化

训练100epoch结果如下，增加epoch还能提升。

3.2 模型预测

在ultralytics-main目录新建一个predict.py脚本，内容如下：

from ultralytics import YOLOif __name__ == '__main__':# Load a modelmodel = YOLO(r"E:\Code\ultralytics-main\runs\detect\train\weights\best.pt")  # load modelmodel.predict(source=r"E:\datasets\visdrone2019\VisDrone2019-DET-test-dev\images\0000006_01111_d_0000003.jpg", save=True, save_conf=True, save_txt=True, name='output')

结果如下：

4 Onnxruntime推理

在ultralytics-main目录新建一个onnx_infer.py脚本，内容如下：
注意：如导出动态onnx，model.export(format=“onnx”, opset=13, dynamic=True)

import argparse
import time 
import cv2
import numpy as npimport onnxruntime as ort  # 使用onnxruntime推理用上，pip install onnxruntime-gpu==1.12.0 -i  https://pypi.tuna.tsinghua.edu.cn/simple，默认安装CPU
import os 
os.environ['CUDA_VISIBLE_DEVICES'] = '0'class YOLOv8:"""YOLOv8 object detection model class for handling inference and visualization."""def __init__(self, onnx_model, imgsz=(640, 640)):"""Initialization.Args:onnx_model (str): Path to the ONNX model."""# 构建onnxruntime推理引擎self.ort_session = ort.InferenceSession(onnx_model,providers=['CUDAExecutionProvider', 'CPUExecutionProvider']if ort.get_device() == 'GPU' else ['CPUExecutionProvider'])print(ort.get_device())# Numpy dtype: support both FP32 and FP16 onnx modelself.ndtype = np.half if self.ort_session.get_inputs()[0].type == 'tensor(float16)' else np.singleself.model_height, self.model_width = imgsz[0], imgsz[1]  # 图像resize大小def __call__(self, im0, conf_threshold=0.4, iou_threshold=0.45):"""The whole pipeline: pre-process -> inference -> post-process.Args:im0 (Numpy.ndarray): original input image.conf_threshold (float): confidence threshold for filtering predictions.iou_threshold (float): iou threshold for NMS.Returns:boxes (List): list of bounding boxes."""# 前处理Pre-processt1 = time.time()im, ratio, (pad_w, pad_h) = self.preprocess(im0)pre_time = round(time.time() - t1, 3)# print('det预处理时间：{:.3f}s'.format(time.time() - t1))# 推理 inferencet2 = time.time()preds = self.ort_session.run(None, {self.ort_session.get_inputs()[0].name: im})[0]# print('det推理时间：{:.2f}s'.format(time.time() - t2))det_time = round(time.time() - t2, 3)# 后处理Post-processt3 = time.time()boxes = self.postprocess(preds,im0=im0,ratio=ratio,pad_w=pad_w,pad_h=pad_h,conf_threshold=conf_threshold,iou_threshold=iou_threshold,)# print('det后处理时间：{:.3f}s'.format(time.time() - t3))post_time = round(time.time() - t3, 3)return boxes, (pre_time, det_time, post_time)# 前处理，包括：resize, pad, HWC to CHW，BGR to RGB，归一化，增加维度CHW -> BCHWdef preprocess(self, img):"""Pre-processes the input image.Args:img (Numpy.ndarray): image about to be processed.Returns:img_process (Numpy.ndarray): image preprocessed for inference.ratio (tuple): width, height ratios in letterbox.pad_w (float): width padding in letterbox.pad_h (float): height padding in letterbox."""# Resize and pad input image using letterbox() (Borrowed from Ultralytics)shape = img.shape[:2]  # original image shapenew_shape = (self.model_height, self.model_width)r = min(new_shape[0] / shape[0], new_shape[1] / shape[1])ratio = r, rnew_unpad = int(round(shape[1] * r)), int(round(shape[0] * r))pad_w, pad_h = (new_shape[1] - new_unpad[0]) / 2, (new_shape[0] - new_unpad[1]) / 2  # wh paddingif shape[::-1] != new_unpad:  # resizeimg = cv2.resize(img, new_unpad, interpolation=cv2.INTER_LINEAR)top, bottom = int(round(pad_h - 0.1)), int(round(pad_h + 0.1))left, right = int(round(pad_w - 0.1)), int(round(pad_w + 0.1))img = cv2.copyMakeBorder(img, top, bottom, left, right, cv2.BORDER_CONSTANT, value=(114, 114, 114))  # 填充# Transforms: HWC to CHW -> BGR to RGB -> div(255) -> contiguous -> add axis(optional)img = np.ascontiguousarray(np.einsum('HWC->CHW', img)[::-1], dtype=self.ndtype) / 255.0img_process = img[None] if len(img.shape) == 3 else imgreturn img_process, ratio, (pad_w, pad_h)# 后处理，包括：阈值过滤与NMSdef postprocess(self, preds, im0, ratio, pad_w, pad_h, conf_threshold, iou_threshold):"""Post-process the prediction.Args:preds (Numpy.ndarray): predictions come from ort.session.run().im0 (Numpy.ndarray): [h, w, c] original input image.ratio (tuple): width, height ratios in letterbox.pad_w (float): width padding in letterbox.pad_h (float): height padding in letterbox.conf_threshold (float): conf threshold.iou_threshold (float): iou threshold.Returns:boxes (List): list of bounding boxes."""x = preds  # outputs: predictions (1, 84, 8400)# Transpose the first output: (Batch_size, xywh_conf_cls, Num_anchors) -> (Batch_size, Num_anchors, xywh_conf_cls)x = np.einsum('bcn->bnc', x)  # (1, 8400, 84)# Predictions filtering by conf-thresholdx = x[np.amax(x[..., 4:], axis=-1) > conf_threshold]# Create a new matrix which merge these(box, score, cls) into one# For more details about `numpy.c_()`: https://numpy.org/doc/1.26/reference/generated/numpy.c_.htmlx = np.c_[x[..., :4], np.amax(x[..., 4:], axis=-1), np.argmax(x[..., 4:], axis=-1)]# NMS filtering# 经过NMS后的值, np.array([[x, y, w, h, conf, cls], ...]), shape=(-1, 4 + 1 + 1)x = x[cv2.dnn.NMSBoxes(x[:, :4], x[:, 4], conf_threshold, iou_threshold)]# 重新缩放边界框，为画图做准备if len(x) > 0:# Bounding boxes format change: cxcywh -> xyxyx[..., [0, 1]] -= x[..., [2, 3]] / 2x[..., [2, 3]] += x[..., [0, 1]]# Rescales bounding boxes from model shape(model_height, model_width) to the shape of original imagex[..., :4] -= [pad_w, pad_h, pad_w, pad_h]x[..., :4] /= min(ratio)# Bounding boxes boundary clampx[..., [0, 2]] = x[:, [0, 2]].clip(0, im0.shape[1])x[..., [1, 3]] = x[:, [1, 3]].clip(0, im0.shape[0])return x[..., :6]  # boxeselse:return []if __name__ == '__main__':# Create an argument parser to handle command-line argumentsparser = argparse.ArgumentParser()parser.add_argument('--det_model', type=str, default=r"E:\Code\ultralytics-main\runs\detect\train\weights\best.onnx", help='Path to ONNX model')parser.add_argument('--source', type=str, default=str(r'E:\datasets\visdrone2019\VisDrone2019-DET-test-dev\images'), help='Path to input image')parser.add_argument('--out_path', type=str, default=str(r'E:\Code\ultralytics-main\runs/res'), help='结果保存文件夹')parser.add_argument('--imgsz_det', type=tuple, default=(640, 640), help='Image input size')parser.add_argument('--classes', type=list, default=['pedestrian', 'people', 'bicycle', 'car', 'van', 'truck', 'tricycle', 'awning-tricycle', 'bus', 'motor'], help='类别')parser.add_argument('--conf', type=float, default=0.25, help='Confidence threshold')parser.add_argument('--iou', type=float, default=0.6, help='NMS IoU threshold')args = parser.parse_args()if not os.path.exists(args.out_path):os.mkdir(args.out_path)print('开始运行：')# Build modeldet_model = YOLOv8(args.det_model, args.imgsz_det)color_palette = np.random.uniform(0, 255, size=(len(args.classes), 3))  # 为每个类别生成调色板for i, img_name in enumerate(os.listdir(args.source)):try:t1 = time.time()# Read image by OpenCVimg = cv2.imread(os.path.join(args.source, img_name))# 检测Inferenceboxes, (pre_time, det_time, post_time) = det_model(img, conf_threshold=args.conf, iou_threshold=args.iou)print('{}/{} ==>总耗时间: {:.3f}s, 其中, 预处理: {:.3f}s, 推理: {:.3f}s, 后处理: {:.3f}s, 识别{}个目标'.format(i+1, len(os.listdir(args.source)), time.time() - t1, pre_time, det_time, post_time, len(boxes)))for (*box, conf, cls_) in boxes:cv2.rectangle(img, (int(box[0]), int(box[1])), (int(box[2]), int(box[3])),color_palette[int(cls_)], 2, cv2.LINE_AA)cv2.putText(img, f'{args.classes[int(cls_)]}: {conf:.3f}', (int(box[0]), int(box[1] - 9)),cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2, cv2.LINE_AA)cv2.imwrite(os.path.join(args.out_path, img_name), img)except Exception as e:print(e)      

资源占用：显存不到2G，RTX4080推理耗时20几毫秒。

结果可视化如下：