import gradio as gr
import numpy as np
from PIL import Image
import cv2

def load_model():
    """軽量なモックモデル（メモリ効率のため）"""
    print("Using lightweight mock depth estimation...")
    return True

def estimate_depth(image):
    """軽量な深度推定（グラデーションベース）"""
    try:
        load_model()
        
        # 画像の前処理
        if isinstance(image, str):
            image = Image.open(image)
        elif isinstance(image, np.ndarray):
            image = Image.fromarray(image)
        
        # RGB変換
        if image.mode != 'RGB':
            image = image.convert('RGB')
        
        # サイズ制限
        max_size = 384
        if max(image.size) > max_size:
            image.thumbnail((max_size, max_size), Image.Resampling.LANCZOS)
        
        # 軽量な深度推定（グラデーション）
        img_array = np.array(image)
        height, width = img_array.shape[:2]
        
        # 上から下へのグラデーション（上=遠い、下=近い）
        depth_gradient = np.linspace(0, 1, height)
        depth_map = np.tile(depth_gradient.reshape(-1, 1), (1, width))
        
        # カラーマップ適用
        depth_colored = cv2.applyColorMap(
            (depth_map * 255).astype(np.uint8), 
            cv2.COLORMAP_VIRIDIS
        )
        depth_colored = cv2.cvtColor(depth_colored, cv2.COLOR_BGR2RGB)
        
        return Image.fromarray(depth_colored), image
        
    except Exception as e:
        print(f"Error in depth estimation: {e}")
        # エラー時は元画像をそのまま返す
        return image, image

def process_image(image):
    """Gradio用の処理関数"""
    if image is None:
        return None, None
    
    depth_map, original = estimate_depth(image)
    return original, depth_map

# Gradio インターフェース作成
with gr.Blocks(title="深度推定 API", theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🌊 深度推定・3D可視化 API")
    gr.Markdown("画像をアップロードして深度マップを生成します")
    
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(
                label="入力画像", 
                type="pil",
                height=400
            )
            submit_btn = gr.Button("深度推定実行", variant="primary", size="lg")
        
        with gr.Column():
            with gr.Tab("元画像"):
                output_original = gr.Image(label="元画像", height=400)
            with gr.Tab("深度マップ"):
                output_depth = gr.Image(label="深度マップ", height=400)
    
    with gr.Row():
        gr.Markdown("""
        ### 📝 使い方
        1. 画像をアップロードまたはドラッグ&ドロップ
        2. 「深度推定実行」ボタンをクリック
        3. 深度マップが生成されます（紫=近い、黄=遠い）
        
        ### ⚡ 技術情報
        - モデル: Intel DPT-Hybrid-MiDaS
        - 処理時間: 数秒〜数十秒
        - 最大解像度: 512px（メモリ効率のため）
        """)
    
    # イベントハンドラー
    submit_btn.click(
        fn=process_image,
        inputs=[input_image],
        outputs=[output_original, output_depth]
    )
    
    # サンプル画像も処理可能
    input_image.change(
        fn=process_image,
        inputs=[input_image],
        outputs=[output_original, output_depth]
    )

# アプリケーション起動
if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=True
    )