import os
import cv2
import numpy as np
import shutil
import random

# ================= 配置 =================
# 原始高清数据集路径
source_root = r"G:\datasets\lolblurtest\test\high_sharp_original"
# 实验根目录
exp_root = r"G:\IR_Experiment\Order_Test"

# 输出路径：生成的“坏图”放这里
target_input_dir = os.path.join(exp_root, "input_dark_noisy")
os.makedirs(target_input_dir, exist_ok=True)

# ================= 降质函数 =================
def degradate_image(img):
    # 1. 变暗 (模拟低光) - 亮度降为 20%
    img = img.astype(np.float32) * 0.2
    
    # 2. 加噪 (模拟高感光度噪声)
    noise = np.random.normal(0, 15, img.shape) # sigma=15
    img = img + noise
    
    # 截断并转回 uint8
    img = np.clip(img, 0, 255).astype(np.uint8)
    return img

# ================= 执行提取与生成 =================
print("🚀 开始构建数据集...")
count = 0

# 遍历源目录
for root, dirs, files in os.walk(source_root):
    # 找到图片文件
    img_files = [f for f in files if f.lower().endswith(('.png', '.jpg', '.jpeg'))]
    
    if img_files:
        # 每个文件夹只取第一张，节省时间，够测就行
        file_name = img_files[0]
        src_path = os.path.join(root, file_name)
        
        # 读取
        img = cv2.imread(src_path)
        if img is None: continue
        
        # 制造降质
        bad_img = degradate_image(img)
        
        # 保存 (重命名以防冲突)
        folder_name = os.path.basename(root)
        save_name = f"{folder_name}_{file_name}"
        save_path = os.path.join(target_input_dir, save_name)
        
        cv2.imwrite(save_path, bad_img)
        count += 1
        print(f"[{count}] 处理: {save_name}")
        
        # 限制数量：为了做 demo，生成 5-10 张就足够说明问题了，多了跑得慢
        if count >= 8: 
            break

print(f"✅ 数据集准备完毕！共 {count} 张。")
print(f"📂 坏图位置: {target_input_dir}")