import gradio as gr
import spaces
import torch
import numpy as np
from PIL import Image
from depth_anything_3.api import DepthAnything3
import tempfile
import imageio.v2 as imageio

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = DepthAnything3.from_pretrained("depth-anything/da3nested-giant-large")
model = model.to(device=device)

@spaces.GPU
def analyze_depth(image, process_res):
    if image is None:
        return None, None, None, None, None

    prediction = model.inference([image], process_res_method='lower_bound_resize', process_res=process_res)

    depth_np = prediction.depth[0].cpu().numpy() if torch.is_tensor(prediction.depth[0]) else prediction.depth[0]
    extrinsics_np = prediction.extrinsics[0].cpu().numpy() if torch.is_tensor(prediction.extrinsics[0]) else prediction.extrinsics[0]
    intrinsics_np = prediction.intrinsics[0].cpu().numpy() if torch.is_tensor(prediction.intrinsics[0]) else prediction.intrinsics[0]

    depth_min = float(depth_np.min())
    depth_max = float(depth_np.max())
    depth_normalized = (depth_max - depth_np) / (depth_max - depth_min)

    depth_8bit_img = Image.fromarray((depth_normalized * 255).astype(np.uint8), mode='L')
    depth_16bit = (depth_normalized * 65535).astype(np.uint16)

    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.png')
    temp_path = temp_file.name
    temp_file.close()
    imageio.imwrite(temp_path, depth_16bit)

    return (
        depth_8bit_img,
        temp_path,
        {"depth_min": depth_min, "depth_max": depth_max, "size": list(depth_np.shape[::-1])},
        {"matrix": extrinsics_np.tolist()},
        {"matrix": intrinsics_np.tolist()}
    )

# 建立 Gradio 介面
demo = gr.Interface(
    fn=analyze_depth,
    inputs=[
        gr.Image(type="pil", label="上傳圖片"),
        gr.Slider(minimum=256, maximum=2048, value=1008, step=16, label="處理解析度 (Process Resolution)")
    ],
    outputs=[
        gr.Image(type="pil", label="8bit 預覽圖"),
        gr.File(label="16bit 深度圖 (PNG)"),
        gr.JSON(label="深度範圍 (Depth Min/Max)"),
        gr.JSON(label="相機外參 (Extrinsics)"),
        gr.JSON(label="相機內參 (Intrinsics)")
    ],
    title="Depth Anything V3 - 深度分析",
    description="上傳圖片，輸出深度圖（白色=近，黑色=遠）及相機參數。16bit PNG 保留更精準的深度值 (0-65535)。"
)

demo.launch()