import gradio as gr
import torch
import numpy as np
import cv2
from PIL import Image
from transformers import pipeline

# Hugging FaceのDepth Anything V2パイプラインをロード
depth_estimator = pipeline(
    task="depth-estimation",
    model="depth-anything/Depth-Anything-V2-Base-hf"
)

def estimate_depth(image):
    # GradioからPIL形式で画像が渡される
    result = depth_estimator(image)
    depth = np.array(result["depth"])
    
    # 深度マップを正規化して見やすく
    depth_normalized = cv2.normalize(depth, None, 0, 255, cv2.NORM_MINMAX)
    depth_colored = cv2.applyColorMap(depth_normalized.astype(np.uint8), cv2.COLORMAP_INFERNO)
    
    return Image.fromarray(depth_colored)

# Gradio UI
demo = gr.Interface(
    fn=estimate_depth,
    inputs=gr.Image(type="pil", label="入力画像をアップロード"),
    outputs=gr.Image(type="pil", label="深度マップ（推定結果）"),
    title="Depth Anything V2 - Web Demo",
    description="LiheYoung/depth-anything-v2-base を使用して画像から深度マップを推定します。"
)

if __name__ == "__main__":
    demo.launch()