| 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!") |
|
|
