import os
import yaml

def generate_txt_file(data_dir, subset, txt_filename):
    subset_dir = os.path.join(data_dir, 'images', subset)
    image_dir = os.path.join(subset_dir+'2017')
    print(image_dir)
    label_dir = os.path.join(data_dir, 'labels', subset + '2017')
    
    image_paths = []
    for filename in os.listdir(image_dir):
        if filename.endswith('.jpg') or filename.endswith('.png'):
            label_filename = filename.replace(".jpg", ".txt")
            txt_path = os.path.join(label_dir, label_filename)
            if not os.path.exists(txt_path):
                continue
            image_path = os.path.join(image_dir, filename)
            image_paths.append(image_path)
    
    txt_path = os.path.join(data_dir, txt_filename)
    with open(txt_path, 'w') as f:
        for image_path in image_paths:
            f.write(image_path + '\n')


def generate_class_txt(coco_dir, yaml_file):
    yaml_path = os.path.join(coco_dir, yaml_file)
    with open(yaml_path, 'r') as f:
        data = yaml.safe_load(f)
    
    class_names = data['names']
    class_txt_path = os.path.join(coco_dir, 'class.txt')
    with open(class_txt_path, 'w') as f:
        for class_name in class_names:
            f.write(class_name + '\n')

def main():
    coco_dir = '/mnt/d/WorkSpace/PyCharmGitWorkspace/ObjectDetectio-watermarking/datasets/coco128'  # 替换为你的 COCO 数据集路径
    # yaml_file = 'coco.yaml'  # COCO YAML 文件名

    # 生成 train.txt
    generate_txt_file(coco_dir, 'train', 'train.txt')
    print("Processed train dataset")

    # 生成 val.txt
    generate_txt_file(coco_dir, 'val', 'val.txt')
    print("Processed val dataset")

    # 生成 test.txt
    # generate_txt_file(coco_dir, 'test', 'test.txt')
    # print("Processed test dataset")

    # 生成 class.txt
    # generate_class_txt(coco_dir, yaml_file)
    # print("Processed class file")

    print("Finished processing COCO dataset")

if __name__ == "__main__":
    main()