人工智能 计算机视觉 yolov9目标检测(训练自己的数据集)

yolov9目标检测(训练自己的数据集)

yolov9目标检测(训练，推理

yolov9目标检测(训练自己的数据集)1. 数据集1.1 数据集1.2 数据集处理

2. 环境配置3. 训练准备流程3.1 新建数据集yaml文件3.2 下载预训练权重3.3 模型文件修改3.4 训练模型

代码地址 论文地址

1. 数据集

1.1 数据集

NWPU VHR-10(Cheng et al.，2016)这个高分辨率(VHR)遥感图像数据集是由西北工业大学(NWPU)构建的，一个具有挑战性的十类地理空间物体检测数据集，用于多类物体检测。这十类对象是飞机、轮船、储油罐、棒球场、网球场、篮球场、田径场、港口、桥梁和车辆。该数据集共包含800幅VHR光学遥感图像，其中715幅彩色图像来自Google Earth，空间分辨率为0.5-2米，85幅全色红外图像来自Vaihingen数据，空间分辨率为0.08米。在该数据集中有两个图像集：正图像集，包括650幅图像，每幅图像包含至少一个待检测的目标;负图像集，包括150幅图像，不包含给定对象类的任何目标。从正面图像集中， 757架飞机，302艘船，655个储罐，390个棒球场，524个网球场，159个篮球场，163个地面田径场，224个港口，124座桥梁和477辆汽车被手动注释了用于地面实况的边界框。负图像集用于基于半监督学习的对象检测和基于弱监督学习的对象检测。

1.2 数据集处理

使用滑窗裁减将图片(标签)裁减成统一大小(320x320，yolo要求图像格式必须为32的倍数)，标签格式处理为类别(0-9)，x(中心，归一化)，y(中心，归一化)，w(归一化)，h(归一化)。处理好后的数据图片放入images文件夹中，label标签放入labels文件夹中。 yolov9数据集格式要求:

dataset/

├─train.txt

├─val.txt

├─test.txt

├─images

│ ├──train

│ │ ├──xxx.jpg

│ │ └──xxx.jpg

│ ├──test

│ │ ├──xxx.jpg

│ │ └──xxx.jpg

│ └──val

│ ├──xxx.jpg

│ └──xxx.jpg

└──labels

├──train

│ ├──xxx.txt

│ └──xxx.txt

├──test

│ ├──xxx.txt

│ └──xxx.txt

└──val

├──xxx.txt

└──xxx.txt

处理数据集代码如下（示例）：

import os

import shutil

from sklearn.model_selection import train_test_split

import copy

# 设置数据文件夹路径

image_dir = "./images"

label_dir = "./labels"

# 获取image和label文件夹中的所有文件名

image_files = os.listdir(image_dir)

label_files = os.listdir(label_dir)

# 确保image和label文件夹中的文件数量相同

assert len(image_files) == len(label_files), "Number of image and label files must be t he same!"

# 将文件名组合为一个列表

label_files=copy.copy(image_files)

for i in range(len(label_files)):

label_files[i]=label_files[i].replace(".png",".txt")

files = list(zip(image_files, label_files))

# 划分数据为训练集和测试集（这里还包括验证集，但你可以根据需要调整比例）

train_files, temp_files = train_test_split(files, test_size=0.4, random_state=42) # 假设40%为测试集

valid_files, test_files = train_test_split(temp_files, test_size=0.5, random_state=42) # 剩下的50%中，再取50%为验证集

print("测试集长度："+str(len(test_files)))

print("训练集长度："+str(len(train_files)))

print("验证集长度："+str(len(valid_files)))

# 创建目录（如果它们不存在）

os.makedirs(image_dir + "/train", exist_ok=True)

os.makedirs(image_dir + "/test", exist_ok=True)

os.makedirs(image_dir + "/val", exist_ok=True)

os.makedirs(label_dir + "/train", exist_ok=True)

os.makedirs(label_dir + "/test", exist_ok=True)

os.makedirs(label_dir + "/val", exist_ok=True)

# 移动文件到相应的目录

for img, lbl in train_files:

shutil.move(os.path.join(image_dir, img), os.path.join(image_dir + "/train", img))

shutil.move(os.path.join(label_dir, lbl), os.path.join(label_dir + "/train", lbl))

for img, lbl in test_files:

shutil.move(os.path.join(image_dir, img), os.path.join(image_dir + "/test", img))

shutil.move(os.path.join(label_dir, lbl), os.path.join(label_dir + "/test", lbl))

for img, lbl in valid_files:

shutil.move(os.path.join(image_dir, img), os.path.join(image_dir + "/val", img))

shutil.move(os.path.join(label_dir, lbl), os.path.join(label_dir + "/val", lbl))

print("Data split completed!")

生成训练集、测试集、验证集txt文件代码如下（示例）：

import os

def write_file_paths_to_txt(directory, output_file):

with open(output_file, 'w') as f:

for root, dirs, files in os.walk(directory):

for file in files:

file_path = os.path.join(root, file)

f.write(file_path + '\n')

# 调用函数，为每个文件夹写入文件路径

write_file_paths_to_txt('./images/train', 'train.txt')

write_file_paths_to_txt('./images/test', 'test.txt')

write_file_paths_to_txt('./images/val', 'val.txt')

2. 环境配置

环境配置：

python：3.8cuda：12.3torch：2.2.1torchaudio: 2.2.1torchvision: 0.17.1ubuntu: 22.04.3gpu: 3060(显存可用5.09G)

其余依赖通过以下命令安装

pip install -r requirements.txt

3. 训练准备流程

3.1 新建数据集yaml文件

在 yolov9/data 路径下新建 vhr10.yaml 文件，以路径下的 coco.yaml 为标准进行参数配置，修改 path 为数据存储路径 修改， names 为对应的标签名。

如下所示:

path: /home/mjy/yolov9/dataset

train: train.txt

val: val.txt

test: test.txt

# Classes

names: ['airplane','ship','storage_tank','baseball_diamond','tennis_court','basketball_court','ground_track_field','harbor','bridge','vehicle']

3.2 下载预训练权重

可下载权重如下所示(提供了-c和-e版本)，下载文件到文件夹下：

ModelTest SizeAPvalAP50valAP75valParam.FLOPsYOLOv9-S64046.8%63.4%50.7%7.2M26.7GYOLOv9-M64051.4%68.1%56.1%20.1M76.8GYOLOv9-C64053.0%70.2%57.8%25.5M102.8GYOLOv9-E64055.6%72.8%60.6%58.1M192.5G

3.3 模型文件修改

在 yolov9/models/detect 路径下选择 yolov9-c.yaml 文件，修改 nc 为训练的标签数量，修改为10

3.4 训练模型

训练模型前需要将util文件下loss_tal.py 文件第167行改为p[0]，不然会报list out of Index错误。 训练模型的参数解释入下

weights：预训练权重路径cfg：模型配置文件路径data：数据配置文件路径epoches：训练 epoch 数量batch-size：训练 batch-size 数量imgsz：图片大小optimizer：优化器

训练时指令设置如下：

python train.py --batch 16 --data data/vhr10.yaml --img 320 --cfg models/detect/yolov9-c.yaml --weights 'yolov9-c.pt' --hyp hyp.scratch-high.yaml --epochs 100

训练结果如下

精彩内容