Yolo txt标签文件转换至Kitti txt标签格式

由于最近用的算法要用到kitti数据格式的标签文件，但是手头的数据呢都是yolo的txt标签文件，于是写了一个脚本来转一下标签文件；（目标检测使用）

• ps：该脚本只是针对一类目标进行标签文件的转换，并且在图片内只有一个被标注的目标。若您的类别数量大于1，请在代码的基础上进行简要修改；

脚本一：

用于检查是否现有标签文件有空的，有空的标签文件的话，就将对应的标签文件以及相应照片删除；

• check_labels.py

# coding: utf-8
# author: hxy
# 2020-6-15
"""
检查标签文件的内容；
删除标件中没有坐标信息的txt文件以及照片文件；
"""
import os
empty_txt = list()


def check_file(labels_folder, images_folder):
    labels = os.listdir(labels_folder)
    for label in labels:
        name = label.split('.')[0]
        size = os.path.getsize(os.path.join(labels_folder, label))
        if size is 0:
        # 确保数据备份，谨慎使用os.remove
            os.remove(os.path.join(labels_folder, name + '.txt'))
            os.remove(os.path.join(images_folder, name + '.jpg'))
            empty_txt.append(name)
    return empty_txt


if __name__ == '__main__':
    empty_list = check_file('./labels',
                            './images')
    print(empty_list)

脚本二：

用于标签格式的转换；

• yolo_txt_to_kitti.py

 coding: utf-8
# author: hxy
# 2020-6-15
"""
代码用于将yolo的txt标签格式转换到kitti的txt标签格式；
"""
import os
import cv2
import time

class_name = ' '


# 将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 + '.jpg'))
            w = img.shape[1]
            h = img.shape[0]
            ori_box = restore_coordinate(label, w, h)
            write_to_kitti_txt('./kitti_labels', 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()
    print('----数据转换开始---')

    restore_results('./images',
                    './labels')

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

互相学习、共同进步！
不足之处欢迎指出！