import os
import shutil

# 定义文件夹路径
base_dir = r"e:\datasets\wlk_right_448"
jpeg_images_dir = os.path.join(base_dir, "JPEGImages")
segmentation_class_dir = os.path.join(base_dir, "SegmentationClass")
annotations_dir = os.path.join(base_dir, "stare_stuct", "annotations")
images_dir = os.path.join(base_dir, "stare_stuct", "images")

# 定义目标文件夹
annotations_training_dir = os.path.join(annotations_dir, "training")
annotations_validation_dir = os.path.join(annotations_dir, "validation")
images_training_dir = os.path.join(images_dir, "training")
images_validation_dir = os.path.join(images_dir, "validation")

# 创建目标文件夹
os.makedirs(annotations_training_dir, exist_ok=True)
os.makedirs(annotations_validation_dir, exist_ok=True)
os.makedirs(images_training_dir, exist_ok=True)
os.makedirs(images_validation_dir, exist_ok=True)

# 读取 train.txt 和 val.txt
train_file = os.path.join(base_dir, "train.txt")
val_file = os.path.join(base_dir, "val.txt")


def read_file_list(file_path):
    with open(file_path, "r") as f:
        return [line.strip() for line in f.readlines()]


train_list = read_file_list(train_file)
val_list = read_file_list(val_file)

# 移动文件函数


def move_files(file_list, src_images_dir, src_labels_dir, dst_images_dir, dst_labels_dir):
    for file_name in file_list:
        # 图片文件
        image_src = os.path.join(src_images_dir, file_name)
        image_dst = os.path.join(dst_images_dir, file_name)
        if os.path.exists(image_src):
            shutil.copy(image_src, image_dst)

        # 标签文件
        label_src = os.path.join(src_labels_dir, file_name)
        label_dst = os.path.join(dst_labels_dir, file_name)
        if os.path.exists(label_src):
            shutil.copy(label_src, label_dst)


# 移动训练集文件
move_files(train_list, jpeg_images_dir, segmentation_class_dir,
           images_training_dir, annotations_training_dir)

# 移动验证集文件
move_files(val_list, jpeg_images_dir, segmentation_class_dir,
           images_validation_dir, annotations_validation_dir)

print("文件组织完成！")