🚀🚀 前言 🚀🚀
改进 YOLO模型的实验数据集选择是非常关键的,因为数据集的质量和多样性会直接影响模型的性能和泛化能力。数据集挑选不好,往往会影响后面实验的顺利程度,甚至可能会导致全局崩盘,从头再来。所以选择好数据集至关重要!
🔥🔥 YOLOv5实验实战篇:
📖 YOLOv5/v7/v8改进实验(一)之数据准备篇
📖 YOLOv5/v7/v8改进实验(二)之数据增强篇
📖 YOLOv5/v7/v8改进实验(三)之训练技巧篇
更新中…
一、数据集选择
在选择改进实验数据集时需要注意的一些关键点:
- 数据集质量和准确性:确保数据集中的标签和边界框信息是准确的。错误的标签会导致模型学习不准确,难以泛化到新数据。
- ⭐数据平衡⭐:尽量保持数据集中各个类别的样本数量平衡。不平衡的数据集可能导致模型对某些类别的性能较差。
- ⭐领域适应⭐:如果您想要做小目标检测相关的实验,那需要选择VisDrone2019等小目标占比较多的公开数据集,确保数据集与目标领域相匹配。
- 数据多样性:确保数据集包含多种不同的目标类别和场景。多样性有助于模型更好地泛化,并适应不同的情况。
- 真实世界数据:使用真实世界的数据集,以反映实际应用场景。合成数据集可以用于预训练,但最终的改进应该在真实数据上进行。
- 数据集分布:了解数据集中物体类别的分布,以便更好地处理不平衡问题。
- 数据集清洗:定期清洗数据集,删除错误的标签和低质量的样本。
- 数据合法性:确保您的数据集采集和使用遵守法律和伦理规定,特别是在涉及隐私和个人信息的情况下。
至于是选择公开数据集还是自制数据集,取决于您自身的条件,如果有条件(无人机等工具)那完全可以自己制作数据集。总之,在改进 YOLO 模型之前,确保你的数据集经过充分的筛选和准备,以获得最佳的结果。
此次实验本人使用的是Kaggle上的一个公开数据集(没条件),总数据量
6048
6048
6048张,有'bicycle', 'bus', 'car', 'motorbike', 'person'
5
5
5个类别。
二、数据数量
- 每类图片数,建议 > 1500张
- 每类实例数,推荐每类标签实例数 > 10000
- 图片多样性,必须代表部署环境,对于现实世界我们推荐图片来自一天中不同时间,不同季节,不同天气,不同光照,不同角度和不同相机等
总结:实际做实验往往达不到以上要求,没有条件则选择3000-6000总量,有条件的可以选择大几千张甚至上万,在图片质量达标的情况下尽可能追求数量,这样训练出来的模型泛化能力更强。
三、数据分布
据集的分布会影响模型的性能和泛化能力。如果训练数据和实际应用场景的数据分布不一致,模型可能会出现过拟合或欠拟合的问题。了解数据分布可以帮助判断是否存在类别不平衡问题,有助于选择适当的算法和模型,以提高模型的泛化能力。
3.1 种类分布
PS:这么一看自己的数据集的分布好像也不是那么的平衡。。。
实现代码
import argparse
import os
from pathlib import Path
import matplotlib.pyplot as plt
import yaml
# 设置中文字体为微软雅黑
plt.rcParams["font.sans-serif"] = "SimHei"
def getClassGtNum(label_dir):
data_dict = {}
assert Path(label_dir).is_dir(), "label_dir is not exist"
txts = os.listdir(label_dir) # 得到label_dir目录下的所有txt GT文件
for txt in txts: # 遍历每一个txt文件
with open(os.path.join(label_dir, txt), "r") as f: # 打开当前txt文件 并读取所有行的数据
lines = f.readlines()
for line in lines: # 遍历当前txt文件中每一行的数据
temp = line.split() # str to list{5}
# print(os.path.join(label_dir, txt))
if temp[0] not in data_dict:
data_dict[temp[0]] = 1
data_dict[temp[0]] += 1
return data_dict
# 画图函数
def plot(SML, classes_names, save=False):
x = classes_names
fig = plt.figure(figsize=(10, 8)) # 画布大小和像素密度
plt.bar(x, SML, width=0.5, align="center", color=["skyblue", "orange", "green"])
for a, b, i in zip(x, SML, range(len(x))): # zip 函数
plt.text(
a, b + 0.01, "%d" % int(SML[i]), ha="center", fontsize=15, color="r"
) # plt.text 函数
plt.xticks(fontsize=15)
plt.yticks(fontsize=15)
plt.xlabel("目标种类", fontsize=16)
plt.ylabel("数量", fontsize=16)
plt.title("数据集各目标种类分布情况", fontsize=16)
# 保存到本地
if save:
plt.savefig("result.png")
plt.show()
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--data", default="data.yaml", help="dataset.yaml path")
parser.add_argument("--labeldir", default="Datasets/labels")
parser.add_argument("--save", action="store_true", help="save final result")
opt = parser.parse_args()
print(opt)
if not isinstance(opt.labeldir, Path):
opt.labeldir = Path(opt.labeldir)
data_dict = getClassGtNum(str(opt.labeldir))
sorted_dict = dict(sorted(data_dict.items()))
with open(opt.data, errors="ignore") as f:
try:
classes_names = yaml.safe_load(f)["names"]
except Exception as e:
print(e)
dict_values = list(sorted_dict.values())
plot(dict_values, classes_names, opt.save)
3.2 GT框分布
实现代码
# 1、统计数据集中小、中、大 GT的个数
# 2、统计某个类别小、中、大 GT的个数
# 3、统计数据集中ss、sm、sl GT的个数
import argparse
import os
from pathlib import Path
import matplotlib.pyplot as plt
# 设置中文字体为微软雅黑
plt.rcParams["font.sans-serif"] = "SimHei"
def getGtAreaAndRatio(label_dir):
"""
得到不同尺度的gt框个数
:params label_dir: label文件地址
:return data_dict: {dict: 3} 3 x {'类别':{'area':[...]}, {'ratio':[...]}}
"""
data_dict = {}
assert Path(label_dir).is_dir(), "label_dir is not exist"
txts = os.listdir(label_dir) # 得到label_dir目录下的所有txt GT文件
for txt in txts: # 遍历每一个txt文件
with open(os.path.join(label_dir, txt), "r") as f: # 打开当前txt文件 并读取所有行的数据
lines = f.readlines()
for line in lines: # 遍历当前txt文件中每一行的数据
temp = line.split() # str to list{5}
coor_list = list(map(lambda x: x, temp[1:])) # [x, y, w, h]
area = float(coor_list[2]) * float(coor_list[3]) # 计算出当前txt文件中每一个gt的面积
# center = (int(coor_list[0] + 0.5*coor_list[2]),
# int(coor_list[1] + 0.5*coor_list[3]))
ratio = round(
float(coor_list[2]) / float(coor_list[3]), 2
) # 计算出当前txt文件中每一个gt的 w/h
# print(os.path.join(label_dir, txt))
if temp[0] not in data_dict:
data_dict[temp[0]] = {}
data_dict[temp[0]]["area"] = []
data_dict[temp[0]]["ratio"] = []
data_dict[temp[0]]["area"].append(area)
data_dict[temp[0]]["ratio"].append(ratio)
return data_dict
def getSMLGtNumByClass(data_dict, class_num):
"""
计算某个类别的小物体、中物体、大物体的个数
params data_dict: {dict: 3} 3 x {'类别':{'area':[...]}, {'ratio':[...]}}
params class_num: 类别 0, 1, 2
return s: 该类别小物体的个数 0 < area <= 0.5%
m: 该类别中物体的个数 0.5% < area <= 1%
l: 该类别大物体的个数 area > 1%
"""
s, m, l = 0, 0, 0
for item in data_dict["{}".format(class_num)]["area"]:
if item <= 0.005:
s += 1
elif item <= 0.01:
m += 1
else:
l += 1
return s, m, l
def getAllSMLGtNum(data_dict, isEachClass=False):
"""
数据集所有类别小、中、大GT分布情况
isEachClass 控制是否按每个类别输出结构
"""
S, M, L = 0, 0, 0
# 需要手动初始化下,有多少个类别就需要写多个
classDict = {
"0": {"S": 0, "M": 0, "L": 0},
"1": {"S": 0, "M": 0, "L": 0},
"2": {"S": 0, "M": 0, "L": 0},
"3": {"S": 0, "M": 0, "L": 0},
"4": {"S": 0, "M": 0, "L": 0},
}
# print(classDict['0']['S'])
# range(class_num)类别数 注意修改!!!
if isEachClass == False:
for i in range(5):
s, m, l = getSMLGtNumByClass(data_dict, i)
S += s
M += m
L += l
return [S, M, L]
else:
for i in range(5):
S = 0
M = 0
L = 0
s, m, l = getSMLGtNumByClass(data_dict, i)
S += s
M += m
L += l
classDict[str(i)]["S"] = S
classDict[str(i)]["M"] = M
classDict[str(i)]["L"] = L
return classDict
# 画图函数
def plotAllSML(SML):
x = ["S:[0, 32x32]", "M:[32x32, 96x96]", "L:[96*96, 640x640]"]
fig = plt.figure(figsize=(10, 8)) # 画布大小和像素密度
plt.bar(x, SML, width=0.5, align="center", color=["skyblue", "orange", "green"])
for a, b, i in zip(x, SML, range(len(x))): # zip 函数
plt.text(
a, b + 0.01, "%d" % int(SML[i]), ha="center", fontsize=15, color="r"
) # plt.text 函数
plt.xticks(fontsize=15)
plt.yticks(fontsize=15)
plt.xlabel("gt大小", fontsize=16)
plt.ylabel("数量", fontsize=16)
plt.title("数据集小、中、大GT分布情况", fontsize=16)
# 保存到本地
plt.savefig("result.png")
plt.show()
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--labeldir", default="Datasets/labels")
opt = parser.parse_args()
print(opt)
if not isinstance(opt.labeldir, Path):
opt.labeldir = Path(opt.labeldir)
data_dict = getGtAreaAndRatio(str(opt.labeldir))
# 1、数据集所有类别小、中、大GT分布情况
# 控制是否按每个类别输出结构
isEachClass = False
SML = getAllSMLGtNum(data_dict, isEachClass)
# print(SML)
if not isEachClass:
plotAllSML(SML)
注意:labels存放标签文件,只需将labeldir
参数指定为labels所在路径即可。
四、数据集划分
划分训练集、验证集和测试集的比例没有一个固定的标准,取决于数据集的大小和可用样本数量,不过要确保数据的随机性,避免数据的偏斜或重复。通常情况下,常见的比例是将数据集划分为训练集、验证集和测试集三部分。
-
训练集是用来训练模型的主要数据集。模型通过训练集学习数据的模式和特征,并调整参数来最小化预测误差。训练集应具有代表性,以涵盖数据的各种变化和情况。
-
验证集用于调整模型的超参数和进行模型选择。超参数是在模型训练之前设置的参数,如学习率、正则化强度等。通过使用验证集,在不同超参数设置下评估模型性能,可以选择最佳的超参数组合,以提高模型的泛化能力。
-
测试集用于评估最终模型的性能。模型在训练和验证期间没有接触到测试集数据,因此测试集提供了一个独立的评估指标,反映了模型在真实场景中的表现。测试集应该是隐藏的,模型在训练过程中不能使用测试集进行调整。
并且比例的选择如下:
- 一般来说,训练集占总数据的60-80%左右,用于模型的训练和参数调整。
- 验证集占总数据的10-20%左右,用于超参数的调整和模型选择。
- 测试集占总数据的10-20%左右,用于最终模型性能的评估。
所以,比例可以是7:2:1,6:2:2,8:1:1等情况。
以YOLOv5训练为例,图像文件存放在images
文件夹中,txt标签文件存放在labels
文件夹中,使用以下代码对其进行划分:
- mydata
- images
- 1.jpg
- 2.jpg
- ...
- labels
- 1.txt
- 2.txt
- ...
import argparse
import glob
from pathlib import Path
import random
import shutil
import os
from tqdm import tqdm
from concurrent.futures import ThreadPoolExecutor
NUM_THREADS = min(8, max(1, os.cpu_count() - 1))
def run(func, this_iter, desc="Processing"):
with ThreadPoolExecutor(max_workers=NUM_THREADS, thread_name_prefix='MyThread') as executor:
results = list(
tqdm(executor.map(func, this_iter), total=len(this_iter), desc=desc)
)
return results
def split_dataset_into_train_val_test(
dataset_dir,
save_dir,
train_ratio=0.7,
val_ratio=0.2,
test_ratio=0.1,
im_suffix='jpg'
):
if isinstance(dataset_dir, str):
dataset_dir = Path(dataset_dir)
image_files = glob.glob(str(dataset_dir / 'images' / f"*.{im_suffix}"))
total_images = len(image_files)
random.shuffle(image_files)
train_split = int(total_images * train_ratio)
val_split = int(total_images * val_ratio)
# test_split = int(total_images * test_ratio)
if train_ratio + val_ratio == 1:
train_images = image_files[:train_split]
val_images = image_files[train_split:]
test_images = []
else:
train_images = image_files[:train_split]
val_images = image_files[train_split : train_split + val_split]
test_images = image_files[train_split + val_split :]
print('*'*25)
print(
"",
f"Total images: {total_images}\n",
f"Train images: {len(train_images)}\n",
f"Val images: {len(val_images)}\n",
f"Test images: {len(test_images)}"
)
print('*'*25)
split_paths = [("train", train_images), ("val", val_images), ("test", test_images)]
for split_name, images in split_paths:
split_dir = Path(save_dir) / split_name
for dir_name in ['images', 'labels']:
if not (split_dir / dir_name).exists():
(split_dir / dir_name).mkdir(exist_ok=True, parents=True)
args_list = [(image, dataset_dir, split_dir) for image in images]
run(process_image, args_list, desc=f"Creating {split_name} dataset")
print(f"Created {split_name} dataset with {len(images)} images.")
def process_image(args):
image_file, dataset_dir, split_dir = args
annotation_file = dataset_dir / 'labels' / f"{Path(image_file).stem}.txt"
assert annotation_file.exists(), f'{annotation_file} 不存在!'
if not has_objects(annotation_file):
return
shutil.copy(image_file, split_dir / "images" / Path(image_file).name)
shutil.copy(annotation_file, split_dir / "labels" / annotation_file.name)
def has_objects(annotation_path):
with open(annotation_path, "r") as f:
lines = f.readlines()
return len(lines) > 0
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument('--data', default='./data') # 数据集Images路径
parser.add_argument('--save', default='./mydata') # 保存路径
parser.add_argument('--images_suffix', default='jpg', help='images suffix') # 图片后缀名
opt = parser.parse_args()
split_dataset_into_train_val_test(
dataset_dir=opt.data,
save_dir=opt.save,
train_ratio=0.7,
val_ratio=0.2,
im_suffix=opt.images_suffix
)