最近在研究DeepStream,其中在学习tlt_cv_examples中的迁移训练例子时,用到的数据集都是kitti数据格式,因此想训练自己的数据的话,需要将自己的YOLO格式的label文件转成kitti数据格式的label,然后在网上看到了一个博主的例子,链接如下:
python实现Yolo txt标签格式转换至Kitti txt标签格式
非常感谢博主贡献的代码,但在实验中发现当一个YOLO的label文件中存在多行bounding box的数据时,始终只能转换一行,还有对一些其他的小问题进行了修改,因此在博主的代码基础上进行了修改。代码如下:
# Yolo txt标签格式转换至Kitti txt标签格式
import os
import cv2
import time
kittiPath = 'kitti_label' # 新生成的kitti数据想存放的文件夹路径,根据需要修改
# 将txt中坐标还原到原始照片的坐标
def restore_coordinate(yolo_bbox, image_w, image_h):
box_w = float(yolo_bbox[3]) * image_w
box_h = float(yolo_bbox[4]) * image_h
x_mid = float(yolo_bbox[1]) * image_w + 1
y_mid = float(yolo_bbox[2]) * image_h + 1
xmin = int(x_mid - box_w / 2)
xmax = int(x_mid + box_w / 2)
ymin = int(y_mid - box_h / 2) + 5 # 增加了一个偏移的量5;
ymax = int(y_mid + box_h / 2) + 5
return [xmin, ymin, xmax, ymax]
# 获取照片的labels文件和images文件,并生成新的标签文件
def restore_results(images_folder, labels_folder):
labels = os.listdir(labels_folder)
for label in labels:
name = label.split('.')[0]
print(name)
with open(os.path.join(labels_folder, label), 'r') as f:
img = cv2.imread(os.path.join(images_folder, name + '.jpg'))
w = img.shape[1]
h = img.shape[0]
info = f.readline()
lines = ""
while info:
label = list(info.split(' '))
ori_box = restore_coordinate(label, w, h)
new_info = label[0] + ' ' + '0.00' + ' ' + '0' + ' ' + '0.00' + ' ' \
+ str(ori_box[0]) + ' ' + str(ori_box[1]) + ' ' + str(ori_box[2]) + ' ' + str(ori_box[3]) \
+ ' ' + '0.00' + ' ' + '0.00' + ' ' + '0.00' + ' ' + '0.00' + ' ' + '0.00' \
+ ' ' + '0.00' + ' ' + '0.00' + '\n'
lines += new_info
info = f.readline()
with open(os.path.join(kittiPath, name + '.txt'), 'w') as fn:
fn.writelines(lines)
fn.close()
# 将转换的坐标值绘制到原始图片上,并显示查看
# cv2.rectangle(img, (ori_box[0], ori_box[1]), (ori_box[2], ori_box[3]), (0, 255, 255), 2)
# cv2.imshow('Transfer_label', img)
# if cv2.waitKey(100) & 0XFF == ord('q'):
# break
f.close()
# cv2.destoryAllWindows()
if __name__ == '__main__':
s = time.time()
imagePath = 'images' # 原本的yolo数据格式的images所在的文件夹,根据自己的修改
labelPath = 'labels' # 原本的yolo数据格式的labels所在的文件夹,根据自己的修改
print('----数据转换开始---')
restore_results(imagePath, labelPath)
print('---耗时:{:.3f}ms'.format(time.time() - s))
print('---数据转换成功---')
转换前的YOLO label文件:
转换成kitti数据格式的label文件: