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

def create_depth_map(image):
    """グラデーション深度マップ生成"""
    if image is None:
        return None, None
    
    try:
        # 画像をnumpy配列に変換
        img_array = np.array(image)
        height, width = img_array.shape[:2]
        
        # 上から下へのグラデーション深度マップ
        depth_map = np.zeros((height, width, 3), dtype=np.uint8)
        
        for y in range(height):
            ratio = y / height
            # 青（遠い）から赤（近い）へのグラデーション
            depth_map[y, :, 0] = int(255 * ratio)      # 赤
            depth_map[y, :, 1] = int(128 * (1 - ratio)) # 緑
            depth_map[y, :, 2] = int(255 * (1 - ratio)) # 青
        
        depth_image = Image.fromarray(depth_map)
        return image, depth_image
        
    except Exception as e:
        print(f"Error: {e}")
        return image, image

# Gradio インターフェース
with gr.Blocks(title="深度推定API") as demo:
    gr.Markdown("# 🌊 深度推定・3D可視化 API")
    gr.Markdown("画像をアップロードして深度マップを生成します")
    
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(
                label="入力画像", 
                type="pil"
            )
        with gr.Column():
            output_original = gr.Image(label="元画像")
            output_depth = gr.Image(label="深度マップ")
    
    # 画像変更時に自動処理
    input_image.change(
        fn=create_depth_map,
        inputs=input_image,
        outputs=[output_original, output_depth]
    )
    
    gr.Markdown("""
    ### 📝 使い方
    - 画像をアップロードすると自動で深度マップが生成されます
    - 青=遠い部分、赤=近い部分
    """)

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