import os
import shutil

if __name__ == '__main__':
    img_src = r"D:\Lab\tensorflow\dataset\Dataset1\VOC\JPEGImages"
    label_src = r'D:\Lab\tensorflow\dataset\Dataset1\VOC\labels'

    train_img = r'D:\Lab\tensorflow\dataset\Dataset1\data_yolo\images\train'
    train_label = r'D:\Lab\tensorflow\dataset\Dataset1\data_yolo\labels\train'

    val_img = r'D:\Lab\tensorflow\dataset\Dataset1\data_yolo\images\val'
    val_label = r'D:\Lab\tensorflow\dataset\Dataset1\data_yolo\labels\val'

    for filename in os.listdir(label_src):
        real_name = filename.split('.')[0]
        if real_name.endswith('8') or real_name.endswith('9') or real_name.endswith('0'):
            # 放入测试集
            print("**正在将{}放入测试集".format(real_name))
            shutil.copy(os.path.join(img_src,real_name+'.jpg'),val_img)
            shutil.copy(os.path.join(label_src , filename), val_label)
        else:
            # 放入训练集
            print("==正在将{}放入训练集".format(real_name))
            shutil.copy(os.path.join(img_src , real_name + '.jpg'), train_img)
            shutil.copy(os.path.join(label_src , filename), train_label)