最近在研究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文件:
在这里插入图片描述

Logo

魔乐社区(Modelers.cn) 是一个中立、公益的人工智能社区,提供人工智能工具、模型、数据的托管、展示与应用协同服务,为人工智能开发及爱好者搭建开放的学习交流平台。社区通过理事会方式运作,由全产业链共同建设、共同运营、共同享有,推动国产AI生态繁荣发展。

更多推荐