折腾了两天YOLOV3,填了几个坑,记录一下。
使用的代码:https://github.com/ultralytics/yolov3
一、数据准备
1.1 训练集和验证集
为了节省数据集的制作时间,直接利用DataFountain“四维图新”自动驾驶感知算法竞赛的数据,如图所示为其中的一个训练集。
下图为官方给出的数据标注,前四个数字为边框的坐标,第五个数字为目标类别,最后一个为置信度
1.2 数据集转化
对于此代码,需要将数据集转化为如下格式
images文件夹下的两个文件夹分别存放训练和验证的图片,labels文件夹下的两个文件夹分别存放训练集和验证集的目标txt文件,其中每一张图片对应一个txt文件。
首先将官方给的train.txt文件分成每张图片对应一个txt文件,代码如下
# 把train.txt中的每张图片的信息提取出来
train_txt_dir = '../train.txt'
train_save_txt = './train_txt/'
f = open(train_txt_dir, 'r')
lines = f.readlines()
for line in lines:
name = line[0:6] + '.txt'
print(name)
txtname = train_save_txt + name
line = line[22:]
if len(line) == 0: # 去掉没有目标的图片
continue
else:
file = open(txtname, 'w')
for content in line.split():
file.write(content[:-2])
file.write('\n')
file.close()
接着将生成的txt文件转化为VOC数据集中的xml格式,代码如下:
import os, sys
import glob
from PIL import Image
img_dir = '/media/tim_mary/study/PycharmProjects/kaggle/PyTorch-YOLOv3/data/adas/images/train'
txt_dir = '/media/tim_mary/study/PycharmProjects/kaggle/PyTorch-YOLOv3/data/adas/data_pro/train_txt'
xml_dir = '/media/tim_mary/study/PycharmProjects/kaggle/PyTorch-YOLOv3/data/adas/data_pro/train_xml'
class_name = ['red', 'green', 'yellow', 'red_left', 'red_right', 'yellow_left',
'yellow_right', 'green_left', 'green_right', 'red_forward', 'green_forward',
'yellow_forward', 'horizon_red', 'horizon_green', 'horizon_yellow', 'off',
'traffic_sign', 'car', 'motor', 'bike', 'bus', 'truck', 'suv', 'express', 'person']
txt_lists = glob.glob(txt_dir + '/*.txt')
img_basenames1 = []
for item in txt_lists:
img_basenames1.append