import os
import json
import numpy as np
from PIL import Image

def convert_scene(scene_path):
    print(f"\n⚙️ 转换场景: {scene_path}")
    rgb_dir = os.path.join(scene_path, "rgb")
    if not os.path.exists(rgb_dir): return
        
    all_files = sorted([f for f in os.listdir(rgb_dir) if f.endswith(".png")])
    if not all_files: return
        
    sample_img_path = os.path.join(rgb_dir, all_files[0])
    with Image.open(sample_img_path) as img:
        w, h = img.size

    with open(os.path.join(scene_path, "intrinsics.txt"), "r") as f:
        line = f.readline().strip()
        focal = float(line.split()[0])
    
    camera_angle_x = 2 * np.arctan(w / (2 * focal))
    pose_dir = os.path.join(scene_path, "pose")
    frames = []
    
    for f in all_files:
        name = os.path.splitext(f)[0]
        pose_path = os.path.join(pose_dir, f"{name}.txt")
        if not os.path.exists(pose_path): continue
        
        c2w = np.loadtxt(pose_path).reshape(4, 4)
        
        # 🔥 核心修复：OpenCV 到 OpenGL 的坐标系转换！
        # 将 Y 轴和 Z 轴反转，让相机“转过头来”看向物体
        c2w[:, 1:3] *= -1 
        
        frames.append({
            "file_path": f"./rgb/{name}", 
            "transform_matrix": c2w.tolist()
        })

    num_train = int(len(frames) * 0.8)
    base_json = {"camera_angle_x": camera_angle_x}
    
    with open(os.path.join(scene_path, "transforms_train.json"), "w") as f:
        json.dump({**base_json, "frames": frames[:num_train]}, f, indent=4)
        
    with open(os.path.join(scene_path, "transforms_test.json"), "w") as f:
        json.dump({**base_json, "frames": frames[num_train:]}, f, indent=4)

    print(f"  ✅ 成功生成修复了坐标系的 transforms_*.json")

target_scenes = ["Chair", "Drums", "Ficus", "Hotdog", "Lego", "Materials", "Mic", "Ship"]
base_path = "/root/autodl-tmp/dataset/Synthetic_NeRF_Verified/Synthetic_NeRF"
for s in target_scenes:
    p = os.path.join(base_path, s)
    if os.path.exists(p):
        convert_scene(p)