Python 代码 由YOLO数据格式转换成Kitti数据集的格式

最近在研究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文件：