一、背景
最近要参加一个比赛,比赛是做图像识别的,比赛提供了部分BDD100K数据集,每个图片是1280*720的交通图片,每个标签中带有物体类别、物体左上锚点、物体右下锚点,如下:
比赛提供的标签数据有9种目标,如下:
比赛前的培训课程上,官方演示了使用YOLO V5模型训练了部分数据,因此我准备使用更高版本的YOLO V10模型训练完整的BDD100K数据集
我的环境信息如下:
- 操作系统windows11
- 显卡为NVIDIA GeForce RTX 3080
- cuda版本11.7
- cudnn版本8.9.7.29
- torch版本2.0.1+cu117
- torchvision版本0.15.0+cu117
二、准备数据集与环境
1. 下载BDD100K数据集
我忘记我的数据集是从哪里下载的了,可以参照BDD100K官方文档去下载,不过我下载的数据集和官网上的路径不大一致,但是差别不大,可以用
我下载下来的数据集路径如下:
images是图片文件路径,里面分为10k和100k,10k的数据体量小,我这里要训练100k的数据,test、train、val分别是测试集、训练集和验证集,里面都是jpg格式的图片数据
labels是标签文件路径,我下载的这里只有100k的,是不是10k的标签在100k里都能找到我没有去验证,对应这images/100k里的路径,labels/100l里有训练集和验证集的标签文件,都是json格式的,文件名与图片一一对应,如下:
标签json文件格式如下:
需要注意的是,BDD100K的数据标签里objects列表后面有一些内容长这样:
我没有去深挖这个标签代表什么,有兴趣可以去研究,因为我在这个项目中用不到这个标签,因此后面做标签转换的时候略过了这些标签。
2. 下载YOLO V10代码
从YOLO V10的Github主页上下载代码,有git的用git,没有git直接下载zip即可
下载完成后,把YOLO v10的代码放到项目路径中,我的路径如图:
3. 标签转换
BDD100K的数据标注方式与YOLO v10的标注方式不同,BDD100K是用左上和右下两个锚点的标注方式,而YOLO V10则是是用的中心点标注方法,例如:
在上图这个例子里,图片中有三个物体,两个人和一条领带,yolo模型的标注文件是txt文件,一行一条标注信息
我们以图片上的齐达内为例,他的标注有5个信息,分别是类别、中心点x相对位置,中心点y相对位置,图像高百分比height,图像宽百分比height
yolo模型不管图片的像素尺寸是多少,图片的长和宽都是0到1的范围,里面的每个物体也是一个矩形,x=0.48表示齐达内的中心点在图像x轴上的位置为0.48,y=0.63表示齐达内的中心点在图像y轴的0.63,height=0.71表示齐达内的高占整个图片高的71%,width=0.69表示齐达内的宽占整个图片宽的69%,应该还是很好理解的。
这种标注相比于BDD100K的标注,好处在于更加通用,如果BDD100K的图像被缩放,不再是1280*720的尺寸,那么所有的标注锚点坐标都需要对应修改
所以我们需要先编写一个python脚本,来吧json格式的BDD100K标注文件转化成txt格式的YOLO v10标注文件,基本思路如下
假设我们有一张图片,大小是1280*720,在图片内部有一个标签为car的物体,左上角锚点的坐标是(100, 120),右下角的锚点坐标是(200, 320),那么:
- 中心点坐标
X = ((200 - 100) / 2 + 100) / 1280 = 150 / 1280 = 0.1171875
- 中心点坐标
Y = ((320 - 120) / 2 + 120) / 720 = 440 / 720 = 0.3055556
- 物体宽度
width = (200 - 100) / 1280 = 0.078125
- 物体高度
height = (320 - 120) / 720 = 0.277778
简化一下,得到公式:
- center_X = (x1 + x2) / 2 / 1280
- center_Y = (y1 + y2) / 2 / 720
- width = (x2 - x1) / 1280
- height = (y2 - y1) / 1280
根据上面的公示编写脚本就可以了,需要注意的是要把训练集的标签和验证集的标签都进行转换,我的转换脚本名为转换标签.py
,路径在cs目录中:
代码如下:
import json
import os
# 输入和输出路径
# input_json_path = '../bdd100k/labels/100k/train'
# output_txt_dir = '../bdd100k/labels/100k/train'
input_json_path = '../bdd100k/labels/100k/val'
output_txt_dir = '../bdd100k/labels/100k/val'
# 类别映射,可以根据需要修改
category_map = {
'bus': 0,
'traffic light': 1,
'traffic sign': 2,
'person': 3,
'bike': 4,
'truck': 5,
'motor': 6,
'car': 7,
'rider': 8,
# 添加其他类别
}
category_keys_list = list(category_map.keys()) # 所有需要识别的物体名称列表
# 确保输出目录存在
os.makedirs(output_txt_dir, exist_ok=True)
def change_label_type(filename):
"""
修改一个文件的标签类别
:param filename: json路径
:return:
"""
with open(filename) as f:
data = json.load(f)
file_name = data.get("name")
frames = data.get("frames")
with open(f"{output_txt_dir}/{file_name}.txt", "w", encoding="utf-8") as out_file:
for each_frame in frames:
objects = each_frame.get("objects")
for each_object in objects:
category = each_object.get("category")
if category not in category_keys_list:
continue
bbox = each_object['box2d']
category_id = category_map.get(category)
xmin, ymin, xmax, ymax = bbox['x1'], bbox['y1'], bbox['x2'], bbox['y2']
x_center = (xmin + xmax) / 2 / 1280
y_center = (ymin + ymax) / 2 / 720
width = (xmax - xmin) / 1280
height = (ymax - ymin) / 720
out_file.write(f"{category_id} {x_center} {y_center} {width} {height}\n")
# 读取BDD100K的标注文件
bdd100k_labels_list = os.listdir(input_json_path)
for label_file in bdd100k_labels_list:
label_file_path = f"{input_json_path}/{label_file}"
change_label_type(label_file_path)
上面的脚本要执行两次,一次转换训练集的标签,一次转换验证集的标签,方法是依次注释或解除注释input_json_path和output_txt_dir这几行代码,input_json_path表示bdd100k的标签路径,output_txt_dir表示转换后的yolo标签输出路径,这里输出路径就直接输出到原标签路径,也就是转换前的json标签和转换后的txt标签在在同一个路径中
脚本中的category_map是参照比赛说明文件编写的(在本文最开头有一个图片表格),因为bdd100k里面用文本而不是数字做为物体类别的标注,为了和比赛一致,所以做了物体名称和数字编号的映射
得到的标签文件如图:
yolo模型的标签文件是txt格式,每行标注一个物体,每行有5个信息,用空格分割,分别是物体类别编号、中心点x坐标、中心点y坐标,物体宽、物体高
4. 开发环境准备
参照YOLO v10 Github主页上的说明,用conda创建一个名为yolov10
的环境,执行下面的代码:
conda create -n yolov10 python=3.9
期间需要输入一个y
然后按回车确认
安装完成后激活环境
conda activate yolov10
随后安装所需依赖,依赖列表文件requirements.txt在yolo v10的源码中,所以我们需要先进入到源码路径
这个步骤就不贴命令了,每个人路径都不一定一样,直接上图
requirements.txt文件就在yolov10-main目录中
执行命令pip install -r requirements.txt
安装依赖,这个过程比较耗时,耐心等待,如果网络不佳,可以尝试切换conda镜像源,方法不在这里赘述,网上很多教程
然后执行pip install -e .
安装yolo命令集,完成开发环境准备
5. GPU环境确认
这里需要注意一点,YOLO v10提供的requiremetns.txt中,torch的版本是2.0.1,这就需要torch版本和CUDA、cuDNN版本对应起来
在torch下载的网站上可以看到,torch-2.0.1版本对应的cuda版本为11.7
上图cu117表示cuda版本为11.7,cp39表示python版本为3.9,win表示windows系统
所以如果你的CUDA不是11.7版本,很有可能无法使用GPU训练,可以使用nvcc --version
显然我的cuda版本是11.2,使用这个版本到了后面也确实无法使用GPU训练,因为torch检测不到GPU设备
FUCK NVIDIA!!
安装CUDA和CUDNN,参照我的另一篇文章,这里不再赘述过程了
还有一点需要注意,那就是YOLO v10的源码包中的requirements.txt里面的torch是cpu版本的,如果要使用GPU,需要去下载GPU版本的torch以及对应版本的torchvision,保险起见仍然使用torch==2.0.1版本,但是需要找到GPU版本
在torch下载的网站上找到的带有cu开头的都是GPU版本,因为我是python3.9环境,CUDA版本为了适配torch=2.0.1升级到了11.7,因此我下载标红的那一个
然后从github主页上查看torch和torchvision的版本对应
因为torch版本是2.0.1,属于2.0系列,所以torchvision的版本需要是0.15,还是在torch下载的网站上找到对应的torchvision版本
安装的时候使用pip install xxxx.whl
安装即可,注意要先安装torchvision,然后安装torch,否则可能会给弄一个版本不对劲的torch!
最后在环境里运行下面这个小脚本,看看能否查看到GPU
import torch
import torchvision
print(torch.__version__) # 打印torch版本
print(torchvision.__version__) # 打印torchvision版本
print(torch.version.cuda) # 打印cuda版本
# 查看是否有可用的GPU
if torch.cuda.is_available():
# 获取GPU的数量
num_gpus = torch.cuda.device_count()
print(f"Available GPUs: {num_gpus}")
# 打印每个GPU的编号和名称
for i in range(num_gpus):
print(f"Device {i}: {torch.cuda.get_device_name(i)}")
else:
print("No GPU available.")
我的输出结果如下:
2.0.1+cu117
0.15.0+cu117
11.7
Available GPUs: 1
Device 0: NVIDIA GeForce RTX 3080
说明torch检测到了我的显卡,我可以用GPU去训练了
有GPU一定用GPU,相比于CPU快的不是一点半点
6. YOLO v10样例
这一步可以不做
在YOLO v10源码路径有个app.py
,通过这是一个简易的web应用,可以上传单个图片进行检测
在YOLO v10源码路径里,执行python app.py
可以启动这个web应用,打开浏览器访问http://127.0.0.1:7860
可以查看web应用的界面,如下:
左上方有个框,我们可以把图片拖进去,然后点击下方的Detect Objects
按钮,等待一会儿就可以对图像中的物体进行识别了
三、使用YOLO v10模型
1. 下载权重文件
首先需要下载YOLO v10的模型,每个模型的下载地址如下,简单理解为越往后越厉害就可以-
- YOLOv10-N:https://github.com/THU-MIG/yolov10/releases/download/v1.1/yolov10n.pt
- YOLOv10-S:https://github.com/THU-MIG/yolov10/releases/download/v1.1/yolov10s.pt
- YOLOv10-M:https://github.com/THU-MIG/yolov10/releases/download/v1.1/yolov10m.pt
- YOLOv10-B:https://github.com/THU-MIG/yolov10/releases/download/v1.1/yolov10b.pt
- YOLOv10-L:https://github.com/THU-MIG/yolov10/releases/download/v1.1/yolov10l.pt
- YOLOv10-X:https://github.com/THU-MIG/yolov10/releases/download/v1.1/yolov10x.pt
下载完成后把.pt文件放置在源码根目录即可,我下载的最后一个 YOLOv10-X
2. 编写yaml配置文件
在yolov10-main中新建一个名为yolov10x.yaml
的文件,内容如下:
path: xxxxxx
train: images/100k/train
val: images/100k/val
test: images/100k/test
names:
0: bus
1: traffic light
2: traffic sign
3: person
4: bike
5: truck
6: motor
7: car
8: rider
上面的yaml文件容易理解,path是数据集的路径,这里我写了绝对路径
train、val、test分别是相对于path的训练集路径、验证集路径以及测试集路径,train和val路径里应当既有图片数据也有标签数据
names是标签类别,我参照比赛要求的表格进行填写,一共九个类别
3. 训练模型
在yolov10-main路径中执行指令即可进行训练:
yolo task=detect mode=train data=yolov10x.yaml model=yolov10x.pt epochs=100 batch=2 device=0 plots=True workers=4
对于上述指令的简要说明如下:
- yolo:运行yolo程序
- task=detect:指定为检测任务
- model=train:指定为训练任务
- data=yolov10x.yaml:依照刚刚编写的yolov10x.yaml文件进行训练
- model=yolov10x.pt:指定下载的yolov10预训练权重文件
- epochs=100:设置训练轮次
- batch=2:设置训练集加载批次,主要是提高训练速度,具体得看你的显卡或者内存容量。如果显存大,则可以设置大一些
- device=0:指定训练设备,如果没有gpu,则令device=cpu,如果有一个gpu,则令device=0,有两个则device=0,1以此类推进行设置
- plots=True:指定在训练过程中生成图表(plots)。这可以帮助你可视化训练进度,如损失函数的变化等
- workers=4:设置进程数,如果显存大,可以往大了设置
运行指令后就开始训练了
估计会训练很久了,训练完了继续更新!
–训练完成后更新–
历时四天,训练了个稀巴烂
内存爆了报MemoryError报了两次,一次是在10/100 Epoch,一次是在40/50 Epoch,人崩溃了许久
每一个Epoch训练都需要接近一小时的时间,40/50这一次明明胜利在望了,看完樊振东vs张本智和的比赛我还过来看了一眼,还在训练,心满意足睡觉去了,没想到啊,看时间应该是刚躺下就报错了,早上起来人很崩溃
我内存有32G,看样子数据量还是太大,如何解决这个问题后面在讨论吧,我现在迫切的想知道训练完之后的步骤,于是又设了一次Epoch=2,只训练两步,历时两个小时顺利训练完成
可以看到训练好的模型有两个,一个是best.pt,一个是last.pt,可以理解成表现最好的模型和最后一步训练出来的模型,他们存放在YOLO v10源码路径的runs/detect/train5/weights目录中,train5是因为这是我第五次训练任务,然后我在train4目录中仍然可以找到best.pt和last.pt,因此得出一个好消息是虽然之前训练因为内存爆了中断,但是我仍然可以使用训练了一半的模型
整个训练完成后会利用10000张验证集图片来检测模型的各个性能指标:Instances, Box, R, mAP50, mAP50-95,对它们的解读是:
- Instances: 验证集中目标实例的数量。
- Box: 预测的边界框的准确性。
- R (Recall): 模型的召回率,表示模型能检测到的目标实例比例。
- mAP50: 在 IoU=0.50 时的平均准确率,反映模型的总体检测精度。
- mAP50-95: 综合 IoU=0.50 到 0.95 的平均准确率,是一个更严格的模型检测性能指标。
详细说明如下(ChatGPT提供):
4. 评估模型
现在我有了train4和train5两次训练的模型,train4是在40/50内存爆了没训练完的模型,train5是只训练了两步的模型,只训练了两步的模型评估效果在上面展示了,各项指标都不理想,于是我准备看一下train4得到的模型效果如何
执行yolo val model=runs/detect/train4/weights/best.pt data=yolov10x.yaml batch=1 imgsz=640
得到的结果如下:
这个比赛任务只关注物体识别的数量,而不关注物体在图片的哪个位置,因此Box这一项不重要,R和mAP50和train5的模型相比都有所提升,但并不是太大,看样子还是得多训练才好
四、使用模型获取目标数量
比赛最终会提供一些新的图片,因此需要调用模型来识别新图片中的物体数量,并且转换成比赛所需的csv格式
1. 预测图片
假设我们最终得到了一个名为best.pt的模型,那么使用这个模型去预测图片可以在YOLO v10源码路径去执行yolo predict model=runs/detect/train4/weights/best.pt source=../../../bdd100k/images/100k/test/ imgsz=640 save_txt=True
这里source是待预测图片的路径,save_txt表示把结果用文本文件的形式进行存储
执行后可以在YOLO v10源码路径的runs/detect/predict
文件夹中找到
跟train过程一样,predict后面的数字表示第几次预测任务,这样方便回顾之前的预测结果
我的在predict2文件夹中,文件夹里的.jpg图片是通过模型标注好的图片,例如:
然后labels文件夹里是文本标注的形式了,打开一个长这样
这个是跟我们前面转换成YOLO的标注方式一样的格式,这个比赛不注重物体位置,只注重物体类别,所以我们只关注每行的第一个数字即可,它对应了一种目标物体
2. 整理比赛所需的csv文件
比赛最终提交一个csv文件,文件有3列,分别表示图片名、人的数量和车的数量,其中人只对应行人,车有5中车
我写了一个脚本,如下
import os
LABELS_PATH = "yolo_source_code/yolov10-main/runs/detect/predict2/labels"
RESULT_CSV_FILENAME = "使用train4的best模型预测结果.csv"
PEOPLE_LABEL_ID = [3]
VEHICLE_LABEL_ID = [0, 4, 5, 6, 7]
file_list = os.listdir(LABELS_PATH)
with open(RESULT_CSV_FILENAME, "w", encoding="utf-8") as f1:
f1.write("image_name,people_num,vehicle_num\n")
for each_label in file_list:
people_num = 0
vehicle_num = 0
with open(f"{LABELS_PATH}/{each_label}", "r", encoding="utf-8") as f2:
image_filename = each_label.replace(".txt", "")
for line in f2:
class_id = int(line.strip().split(" ")[0])
if class_id in PEOPLE_LABEL_ID:
people_num += 1
elif class_id in VEHICLE_LABEL_ID:
vehicle_num += 1
f1.write(f"{image_filename},{people_num},{vehicle_num}\n")
只需设置LABELS_PATH和RESULT_CSV_FILENAME这两个变量值即可,第一个是预测后labels的路径,第二个是输出的csv文件名
执行脚本得到的csv文件格式如图
3. 检查模型在验证集、训练集上的准确性
可以用类似的方法检查模型在验证、训练集上的准确性,具体做法是用模型先预测一下训练集、验证集的图片,整理出csv文件
然后再把原本训练集、验证集的标注数据转化成一样的csv,然后看看完全相同的行有多少,看看准确率的百分比如何
五、写在最后
我不确定这个模型的最终效果如何,由于比赛并不透明公开成绩和答案,因此这个谁也说不准,重要的是在这个过程中学习知识吧
后续我准备想办法解决一下内存爆炸的问题,不行就加内存条看看行不行