PART1/DarkIR/prepare_darkir_data.py

import os
import shutil
import glob

# ================= 配置区域 =================
# 1. 原始数据集根目录 (包含 0060, 0061 等场景文件夹)
source_root = r"G:\datasets\lolblurtest\test\low_blur_noise"

# 2. 目标测试输入目录 (DarkIR 将读取这个文件夹)
target_dir = r"G:\IR_Experiment\DarkIR\test_input"

# ================= 执行逻辑 =================
def prepare_dataset():
    if not os.path.exists(source_root):
        print(f"❌ 错误：源目录不存在 -> {source_root}")
        return

    # 如果目标目录存在，先清空，防止旧数据干扰
    if os.path.exists(target_dir):
        shutil.rmtree(target_dir)
    os.makedirs(target_dir)

    print(f"🚀 开始从 {source_root} 抽取数据...")
    
    # 遍历源目录下的所有子文件夹
    scene_folders = [f for f in os.listdir(source_root) if os.path.isdir(os.path.join(source_root, f))]
    count = 0

    for scene in scene_folders:
        scene_path = os.path.join(source_root, scene)
        
        # 查找图片文件 (png, jpg)
        images = [f for f in os.listdir(scene_path) if f.lower().endswith(('.png', '.jpg', '.jpeg'))]
        
        if images:
            # 取第一张图片
            img_name = images[0]
            src_file = os.path.join(scene_path, img_name)
            
            # 重命名保存，带上场景号，方便区分 (例如: 0060_001.png)
            new_name = f"{scene}_{img_name}"
            dst_file = os.path.join(target_dir, new_name)
            
            shutil.copy2(src_file, dst_file)
            print(f"   [提取] 场景 {scene}: {img_name} -> {new_name}")
            count += 1
    
    print(f"\n✅ 数据准备完毕！共抽取 {count} 张图片。")
    print(f"📂 输入文件夹: {target_dir}")

if __name__ == "__main__":
    prepare_dataset()