Spaces:
Runtime error
Runtime error
File size: 4,583 Bytes
2e78bab |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 |
import gradio as gr
import torch
import numpy as np
from PIL import Image
import io
from transformers import DPTImageProcessor, DPTForDepthEstimation
import cv2
# グローバル変数でモデルを保持
processor = None
model = None
def load_model():
"""モデルを一度だけ読み込む"""
global processor, model
if processor is None or model is None:
print("Loading depth estimation model...")
processor = DPTImageProcessor.from_pretrained("Intel/dpt-hybrid-midas")
model = DPTForDepthEstimation.from_pretrained("Intel/dpt-hybrid-midas")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)
model.eval()
print(f"Model loaded on {device}")
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 = 512
if max(image.size) > max_size:
image.thumbnail((max_size, max_size), Image.Resampling.LANCZOS)
# 推論実行
inputs = processor(images=image, return_tensors="pt")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
inputs = {k: v.to(device) for k, v in inputs.items()}
with torch.no_grad():
outputs = model(**inputs)
predicted_depth = outputs.predicted_depth
# 深度マップの後処理
depth = predicted_depth.squeeze().cpu().numpy()
depth_min = depth.min()
depth_max = depth.max()
if depth_max - depth_min > 0:
depth_normalized = (depth - depth_min) / (depth_max - depth_min)
else:
depth_normalized = np.zeros_like(depth)
# カラーマップ適用
depth_colored = cv2.applyColorMap(
(depth_normalized * 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
) |