Generate JDE dot train

Dot train format for JDE

• 生成train格式文件
• 根据数据集的格式自己筛选

from glob2 import glob
import os

images = r"C:\Users\lpf\Desktop\MOT\Code-MOT\Towards-Realtime-MOT-master\data\KITTI_Det\images"
dot_train = r"C:\Users\lpf\Desktop\MOT\Code-MOT\Towards-Realtime-MOT-master\data\kitti_det.train"

path_labels = r"C:\Users\lpf\Desktop\MOT\Code-MOT\Towards-Realtime-MOT-master\data\KITTI_Det\labels_with_ids"


def write_dot_train(path_dot_train, img_path_list):
    with open(path_dot_train, "w") as f:
        for path in img_path_list:
            f.writelines("{}\n".format(path))

# 本方法为detection数据集转换专用
def generate_dot_train(path_labels):
    include_object = []
    include_img = []
    labels_txt = glob(path_labels + "/*.txt")
    for i in labels_txt:
        file_size = os.path.getsize(i)
        if file_size > 0:
            include_object.append(i)
            # print("File size is :", file_size, "bytes")
            # print(i)
            a = os.path.split(i)[-1]
            temp = a[:-4]
            temp = temp + ".png"
            # temp = os.path.join(images, temp)
            temp = "KITTI_Det/images/{}".format(temp)
            include_img.append(temp)
            # break
    return include_img
        # pass


include_image = generate_dot_train(path_labels=path_labels)
write_dot_train(dot_train,include_image)

# print(len(labels_txt))
# print(len(include_object))

# 本方法为 tracking专用
def generate_kitti_train():
    # 注意本处生成的是以image为基准的train文件，包含所有图片，而通常情况下筛选后的lebels_with_ids则不行，会偏少
    # 修改就是使用 lebels_with_ids 为基准，然后将txt 使用replace 方法改为png
    path = os.walk(images)
    for root, directories, files in path:
        for _dir in directories:
            line = "/KITTI_Det/images/{}/".format(_dir)
            _dir = os.path.join(r"E:\8_DataSet\KITTI_tracking\image_02", _dir)
            # dir_list.append(_dir)
            txt_list = os.listdir(_dir)
            with open("kitti-img.train", 'a') as f:
                for item in txt_list:
                    line1 = line + item
                    f.writelines(line1)
                    f.writelines("\n")