将KITTI中的2D检测数据集转换为YOLOV5能训练的格式
下载KITTI数据集中的2d检测图片
关于该数据集的介绍，参考: kitti目标检测2D数据集
注意：
- 2d数据集的连接
- 下载数据需要自己注册一个kitti的账号
- 下数据集和标注文件（详见该部分的参考）
去除DontCare标注
DontCare 标签表示物体未被标记的区域，例如因为它们离激光扫描仪太远。为防止此类对象干扰检测结果，我们直接忽略数据集中的 DontCare 区域。此部分及其后续部分参考：用yolov5训练kitti数据集，并在该博客的基础上加以修改。
代码实现：
import glob
import os
import string
def show_category(txt_list):
category_list = []
for item in txt_list:
try:
with open(item) as tdf:
for each_line in tdf:
labeldata = each_line.strip().split(' ') # 去掉前后多余的字符并把其分开
category_list.append(labeldata[0]) # 只要第一个字段，即类别
except IOError as ioerr:
print('File error:'+str(ioerr))
unique_category = set(category_list)
return unique_category
def process_labels(txt_list):
for item in txt_list:
new_txt = []
new_txtname = str(item).replace('label_2', 'new_label')
try:
with open(item, 'r') as r_tdf:
res = r_tdf.readlines()
# print(res)
for row in res:
category = str(row).split(' ')[0]
# print(category)
if category == 'DontCare': # 忽略Dontcare类
continue
else:
new_txt.append(row) # 将其余类别重新写入新的txt文件
# print(new_txt)
if not os.path.exists('new_label'):
os.mkdir('new_label')
with open(new_txtname, 'w+') as w_tdf: # w+是打开原文件将内容删除，另写新内容进去
for temp in new_txt:
w_tdf.write(str(temp))
except IOError as ioerr:
print('File error:' + str(ioerr))
if __name__ == '__main__':
txt_list = glob.glob('label_2/*.txt')
# classes = show_category(txt_list)
# print(classes) # 输出集合
# txt_list = txt_list[:10]
process_labels(txt_list)
转换为XML文件
原数据集的标注格式并不适合在yolov5中使用。故在此，我们先将其转换为xml格式的标注，之后再转换为归一化坐标的txt格式标注。
代码实现：
# kitti_txt_to_xml.py
# encoding:utf-8
# 根据一个给定的XML Schema，使用DOM树的形式从空白文件生成一个XML
from xml.dom.minidom import Document
import cv2
import os
def generate_xml(name, split_lines, img_size, class_ind):
doc = Document() # 创建DOM文档对象
annotation = doc.createElement('annotation')
doc.appendChild(annotation)
title = doc.createElement('folder')
title_text = doc.createTextNode('KITTI')
title.appendChild(title_text)
annotation.appendChild(title)
img_name=name+'.png'
title = doc.createElement('filename')
title_text = doc.createTextNode(img_name)
title.appendChild(title_text)
annotation.appendChild(title)
source = doc.createElement('source')
annotation.appendChild(source)
title = doc.createElement('database')
title_text = doc.createTextNode('The KITTI Database')
title.appendChild(title_text)
source.appendChild(title)
title = doc.createElement('annotation')
title_text = doc.createTextNode('KITTI')
title.appendChild(title_text)
source.appendChild(title)
size = doc.createElement('size')
annotation.appendChild(size)
title = doc.createElement('width')
title_text = doc.createTextNode(str(img_size[1]))
title.appendChild(title_text)
size.appendChild(title)
title = doc.createElement('height')
title_text = doc.createTextNode(str(img_size[0]))
title.appendChild(title_text)
size.appendChild(title)
title = doc.createElement('depth')
title_text = doc.createTextNode(str(img_size[2]))
title.appendChild(title_text)
size.appendChild(title)
for split_line in split_lines:
line=split_line.strip().split()
if line[0] in class_ind:
object = doc.createElement('object')
annotation.appendChild(object)
title = doc.createElement('name')
title_text = doc.createTextNode(line[0])
title.appendChild(title_text)
object.appendChild(title)
bndbox = doc.createElement('bndbox')
object.appendChild(bndbox)
title = doc.createElement('xmin')
title_text = doc.createTextNode(str(int(float(line[4]))))
title.appendChild(title_text)
bndbox.appendChild(title)
title = doc.createElement('ymin')
title_text = doc.createTextNode(str(int(float(line[5]))))
title.appendChild(title_text)
bndbox.appendChild(title)
title = doc.createElement('xmax')
title_text = doc.createTextNode(str(int(float(line[6]))))
title.appendChild(title_text)
bndbox.appendChild(title)
title = doc.createElement('ymax')
title_text = doc.createTextNode(str(int(float(line[7]))))
title.appendChild(title_text)
bndbox.appendChild(title)
# 将DOM对象doc写入文件
if not os.path.exists('Annotations'):
os.mkdir('Annotations')
f = open('Annotations/'+name+'.xml', 'w')
f.write(doc.toprettyxml(indent = ''))
f.close()
if __name__ == '__main__':
class_ind = ('Car', 'Van', 'Truck', 'Pedestrian', 'Person_sitting', 'Cyclist', 'Tram', 'Misc')
cur_dir = os.getcwd()
labels_dir = os.path.join(cur_dir, 'new_label')
for parent, dirnames, filenames in os.walk(labels_dir): # 分别得到根目录，子目录和根目录下文件
for file_name in filenames:
full_path = os.path.join(parent, file_name) # 获取文件全路径
f = open(full_path)
split_lines = f.readlines()
name = file_name[:-4] # 后四位是扩展名.txt，只取前面的文件名
img_name = name+'.png'
img_path = os.path.join('/home/dcr/kitti/training/image_2', img_name) # 路径需要自行修改
img_size = cv2.imread(img_path).shape
generate_xml(name, split_lines, img_size, class_ind)
print('all txts has converted into xmls')
xml转换为txt并划分标签的训练集和验证集
此部分先将标签转换为txt格式，然后划分为训练集和验证集，数据集一共7481张图片，其中划分出验证集1500张图片（测试集在下载的原始数据解压之后就有，不需要单独划分）。
代码实现：
# xml_to_yolo_txt.py
import glob
import os
import random
import xml.etree.ElementTree as ET
import numpy as np
# 转换一个xml文件为txt
def single_xml_to_txt(xml_file, save_path: str):
tree = ET.parse(xml_file)
root = tree.getroot()
print(xml_file)
txt_file = str(xml_file.split('.')[1]).replace('/Annotations', save_path) + '.txt'
with open(txt_file, 'w') as txt_file:
for member in root.findall('object'):
#filename = root.find('filename').text
picture_width = int(root.find('size')[0].text)
picture_height = int(root.find('size')[1].text)
class_name = member[0].text
# 类名对应的index
class_num = class_names.index(class_name)
box_x_min = int(member[1][0].text) # 左上角横坐标
box_y_min = int(member[1][1].text) # 左上角纵坐标
box_x_max = int(member[1][2].text) # 右下角横坐标
box_y_max = int(member[1][3].text) # 右下角纵坐标
print(box_x_max, box_x_min, box_y_max, box_y_min)
# 转成相对位置和宽高
x_center = float(box_x_min + box_x_max) / (2 * picture_width)
y_center = float(box_y_min + box_y_max) / (2 * picture_height)
width = float(box_x_max - box_x_min) / picture_width
height = float(box_y_max - box_y_min) / picture_height
print(class_num, x_center, y_center, width, height)
txt_file.write(str(class_num) + ' ' + str(x_center) + ' ' + str(y_center) + ' ' + str(width) + ' ' + str(height) + '\n')
# 转换文件夹下的所有xml文件为txt
def dir_xml_to_txt(path):
# 保存的txt文件路径
if not os.path.exists('labels/train'):
os.makedirs('labels/train') # 递归创建
if not os.path.exists('labels/val'):
os.makedirs('labels/val')
xml_list = glob.glob(path + '*.xml')
random.shuffle(xml_list)
xml_val = xml_list[:1500] # 验证数据
xml_train = xml_list[1500:] # 训练数据
for xml_file in xml_val:
single_xml_to_txt(xml_file,'labels/val')
for xml_file in xml_train:
single_xml_to_txt(xml_file,'labels/train')
if __name__ == '__main__':
# 这里的类名为我们xml里面的类名，顺序现在不需要考虑
class_names = ['Car', 'Van', 'Truck', 'Pedestrian', 'Person_sitting', 'Cyclist', 'Tram', 'Misc']
# xml文件路径
path = './Annotations/'
dir_xml_to_txt(path)
print("Done")
根据标签的划分结果将图片分开
将图片从原始数据中分别复制到指定的文件夹。
代码实现：
import glob
import os
import shutil
import re
def img_split(origin_path, target_path, label_path):
if not os.path.exists(target_path):
os.makedirs(target_path)
origin_img_path = origin_path # 原图片路径 'training/image_2/'
train_list = glob.glob(label_path + '*.txt') # 标签文件的路径'label_train/*.txt’
# print(train_list)
pattern = re.compile(r'\d+') # 使用正则表达式得到标签中的文件名（6个数字）
img_list = pattern.findall(str(train_list))
for img in img_list:
img_file = origin_img_path + img + '.png'
img_path = target_path + img + '.png' # 目标文件路径'img_train/*.png'
# print(img_file)
shutil.copy(img_file, img_path) # 从原路径复制图片到目标路径
if __name__ == '__main__':
# 复制训练集图片
img_split('/home/dcr/kitti/training/image_2/', 'images/train/', 'labels/train/')
# 复制验证集图片
img_split('/home/dcr/kitti/training/image_2/', 'images/val/', 'labels/val/')
数据集处理完成
按顺序执行完上述步骤之后，其中的images和labels文件夹就可以直接在 yolov5 中使用。整个过程中主要参考：
- https://zhuanlan.zhihu.com/p/57493587
- https://blog.csdn.net/qq_45978858/article/details/119718516
1757
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
