python实现Yolo txt标签格式转换至Kitti txt标签格式

最新推荐文章于 2024-04-21 17:49:50 发布

南沙的星星

最新推荐文章于 2024-04-21 17:49:50 发布

阅读量589

点赞数

本文链接：https://blog.csdn.net/weixin_41783910/article/details/108228228

版权

# Yolo txt标签格式转换至Kitti txt标签格式
# author: zhou jinxing
import os
import cv2
import time

class_name = ' '
kittiPath = '/workspace/tlt-experiments/data_fire/data4test/training/kitti_label/'


# 将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]


# 生成kitti格式的txt标签文件
def write_to_kitti_txt(save_path, box, name):
    with open(os.path.join(save_path, name + '.txt'), 'w') as f:
       # kitti标签文件中内容共有15个参数
        new_info = class_name + ' ' + '0.00' + ' ' + '0' + ' ' + '0.00' + ' '\
                   + str(box[0]) + ' ' + str(box[1]) + ' ' + str(box[2]) + ' ' + str(box[3]) \
                   + ' ' + '0.00' + ' ' + '0.00' + ' ' + '0.00' + ' ' + '0.00' + ' ' + '0.00' \
                   + ' ' + '0.00' + ' ' + '0.00'
        f.writelines(new_info)
    f.close()


# 获取照片的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:
            info = f.readline().strip('\n')
            label = list(info.split(' '))
            img = cv2.imread(os.path.join(images_folder, name + '.png'))
            w = img.shape[1]
            h = img.shape[0]
            ori_box = restore_coordinate(label, w, h)
            #write_to_kitti_txt('./kitti_labels', ori_box, name)
            write_to_kitti_txt(kittiPath, ori_box, name)
            
            
            # 将转换的坐标值绘制到原始图片上，并显示查看
            # 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 = '/workspace/tlt-experiments/data_fire/data4test/training/image_2/'
    labelPath = '/workspace/tlt-experiments/data_fire/data4test/training/label_2/'
    print('----数据转换开始---')

    restore_results(imagePath,labelPath)

    print('---耗时：{:.3f}ms'.format(time.time() - s))
    print('---数据转换成功---')



----数据转换开始---
000000
000001
000002
000003
000004
000005
000006
000007
000008
000009
---耗时：0.102ms
---数据转换成功---

参考：https://blog.csdn.net/weixin_38106878/article/details/106782995?biz_id=102&utm_term=Yolov%20txt%20%E6%A0%87%E7%AD%BE%E6%96%87%E4%BB%B6%E8%BD%AC%E6%8D%A2%E8%87%B3Kitti&utm_medium=distribute.pc_search_result.none-task-blog-2_allsobaiduweb~default-0-106782995&spm=1018.2118.3001.4187