import os
from PIL import Image
from concurrent.futures import ThreadPoolExecutor
from tqdm import tqdm

DATASETS = [
    "/root/autodl-tmp/dataset/360",
    "/root/autodl-tmp/dataset/tnt"
]

def process_image(src_path, dst_path, scale):
    if os.path.exists(dst_path):
        return
    try:
        with Image.open(src_path) as img:
            new_size = (int(img.width / scale), int(img.height / scale))
            resized_img = img.resize(new_size, Image.Resampling.LANCZOS)
            resized_img.save(dst_path)
    except Exception as e:
        print(f"Error processing {src_path}: {e}")

def downsample_scene(scene_dir):
    img_dir = os.path.join(scene_dir, "images")
    if not os.path.exists(img_dir):
        return

    images = [f for f in os.listdir(img_dir) if f.lower().endswith(('.png', '.jpg', '.jpeg'))]
    
    # 需要生成的降采样层级
    scales = [2, 4] 
    
    tasks = []
    for scale in scales:
        target_dir = os.path.join(scene_dir, f"images_{scale}")
        os.makedirs(target_dir, exist_ok=True)
        
        for img_name in images:
            src_path = os.path.join(img_dir, img_name)
            dst_path = os.path.join(target_dir, img_name)
            tasks.append((src_path, dst_path, scale))
            
    return tasks

if __name__ == "__main__":
    all_tasks = []
    print("[Preprocessing] Scanning datasets for missing resolution folders...")
    for ds_root in DATASETS:
        if not os.path.exists(ds_root): continue
        for scene in os.listdir(ds_root):
            scene_dir = os.path.join(ds_root, scene)
            if os.path.isdir(scene_dir):
                tasks = downsample_scene(scene_dir)
                if tasks: all_tasks.extend(tasks)
                
    if not all_tasks:
        print("[Preprocessing] All resolution folders exist. Nothing to do.")
    else:
        print(f"[Preprocessing] Generating {len(all_tasks)} downsampled images. This will use all CPU cores...")
        with ThreadPoolExecutor() as executor:
            list(tqdm(executor.map(lambda p: process_image(*p), all_tasks), total=len(all_tasks)))
        print("[Preprocessing] Done!")