把train数据集生成txt(test同理)

import cv2
import numpy as np
import os
import sys
import pickle

data_dir = os.path.join("./")
train_o_dir = os.path.join(data_dir,"train")

def my_mkdir(my_dir):
    if not os.path.isdir(my_dir):
        os.makedirs(my_dir)

# 生成训练集图片，
if __name__ == '__main__':
    data_path = os.path.join(data_dir, "TrainJPEGImages")
    print(data_path + " is loading...")

    datatraint = os.path.join(data_dir, "train.txt")
    ft = open(datatraint, 'w')
    print(data_path)
    img_list = os.listdir(data_path)
    print(img_list)
    for i in range(len(img_list)):
        if not img_list[i].endswith('jpg'):         # 若不是jpg文件，跳过
            continue
        img_path = os.path.join(data_path, img_list[i])
        line = img_path + '\n'
        ft.write(line)
    ft.close()   

    print(data_path + " loaded.")