TK156 commited on
Commit
17b7f25
·
verified ·
1 Parent(s): 8217c26

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +59 -79
app.py CHANGED
@@ -6,87 +6,67 @@
6
  from PIL import Image
7
  import cv2
8
 
9
- class DepthEstimationAPI:
10
- def __init__(self):
11
- self.device = "cpu" # Hugging Face Spacesは無料版でCPUのみ
12
- print(f"Using device: {self.device}")
13
-
14
- model_name = "depth-anything/Depth-Anything-V2-Small-hf"
15
- self.processor = AutoImageProcessor.from_pretrained(model_name)
16
- self.model =
17
- AutoModelForDepthEstimation.from_pretrained(model_name)
18
- self.model.to(self.device)
19
- self.model.eval()
20
- print("Model loaded successfully")
21
-
22
- def predict(self, image_input):
23
- """Process image and return depth map"""
24
- if image_input is None:
25
- return None, None
26
-
27
- try:
28
- # PILイメージに変換
29
- if hasattr(image_input, 'convert'):
30
- image = image_input.convert('RGB')
31
- else:
32
- image = Image.open(image_input).convert('RGB')
33
-
34
- # サイズ調整(メモリ節約)
35
- max_size = 256
36
- if max(image.size) > max_size:
37
- ratio = max_size / max(image.size)
38
- new_size = tuple(int(dim * ratio) for dim in
39
- image.size)
40
- image = image.resize(new_size,
41
- Image.Resampling.LANCZOS)
42
-
43
- # 深度推定
44
- inputs = self.processor(images=image, return_tensors="pt")
45
-
46
- with torch.no_grad():
47
- outputs = self.model(**inputs)
48
- depth = outputs.predicted_depth.squeeze().cpu().numpy()
49
-
50
- # 深度マップ可視化
51
- depth_normalized = ((depth - depth.min()) / (depth.max() -
52
  depth.min()) * 255).astype(np.uint8)
53
- depth_colored = cv2.applyColorMap(depth_normalized,
54
  cv2.COLORMAP_VIRIDIS)
55
- depth_colored = cv2.cvtColor(depth_colored,
56
- cv2.COLOR_BGR2RGB)
57
- depth_image = Image.fromarray(depth_colored)
58
-
59
- return image, depth_image
60
-
61
- except Exception as e:
62
- print(f"Error in prediction: {e}")
63
- return image_input, None
64
-
65
- # APIインスタンス初期化
66
- api = DepthEstimationAPI()
67
-
68
- # Gradioインターフェース作成
69
- with gr.Blocks(title="Depth Estimation API") as demo:
70
- gr.Markdown("# 深度推定 API")
71
- gr.Markdown("DepthAnything V2を使用したAI深度推定")
72
-
73
- with gr.Row():
74
- with gr.Column():
75
- input_image = gr.Image(type="pil",
76
- label="画像をアップロード")
77
- submit_btn = gr.Button("深度マップ生成", variant="primary")
78
-
79
- with gr.Column():
80
- output_original = gr.Image(type="pil", label="元画像")
81
- output_depth = gr.Image(type="pil", label="深度マップ")
82
-
83
- # ボタンクリックで処理実行
84
- submit_btn.click(
85
- fn=api.predict,
86
- inputs=[input_image],
87
- outputs=[output_original, output_depth]
88
- )
89
 
90
- # Hugging Face Spaces用起動設定
91
  if __name__ == "__main__":
92
  demo.launch()
 
6
  from PIL import Image
7
  import cv2
8
 
9
+ # デバイス設定
10
+ device = "cpu"
11
+ print(f"Using device: {device}")
12
+
13
+ # モデル読み込み
14
+ model_name = "depth-anything/Depth-Anything-V2-Small-hf"
15
+ processor = AutoImageProcessor.from_pretrained(model_name)
16
+ model = AutoModelForDepthEstimation.from_pretrained(model_name)
17
+ model.to(device)
18
+ model.eval()
19
+ print("Model loaded successfully")
20
+
21
+ def predict_depth(image):
22
+ """深度推定関数"""
23
+ if image is None:
24
+ return None, None
25
+
26
+ try:
27
+ # 画像処理
28
+ if hasattr(image, 'convert'):
29
+ image = image.convert('RGB')
30
+
31
+ # サイズ調整
32
+ max_size = 256
33
+ if max(image.size) > max_size:
34
+ ratio = max_size / max(image.size)
35
+ new_size = tuple(int(dim * ratio) for dim in image.size)
36
+ image = image.resize(new_size, Image.Resampling.LANCZOS)
37
+
38
+ # 深度推定
39
+ inputs = processor(images=image, return_tensors="pt")
40
+
41
+ with torch.no_grad():
42
+ outputs = model(**inputs)
43
+ depth = outputs.predicted_depth.squeeze().cpu().numpy()
44
+
45
+ # 可視化
46
+ depth_norm = ((depth - depth.min()) / (depth.max() -
 
 
 
 
 
47
  depth.min()) * 255).astype(np.uint8)
48
+ depth_colored = cv2.applyColorMap(depth_norm,
49
  cv2.COLORMAP_VIRIDIS)
50
+ depth_colored = cv2.cvtColor(depth_colored, cv2.COLOR_BGR2RGB)
51
+ depth_image = Image.fromarray(depth_colored)
52
+
53
+ return image, depth_image
54
+
55
+ except Exception as e:
56
+ print(f"Error: {e}")
57
+ return image, None
58
+
59
+ # Gradioインターフェース
60
+ demo = gr.Interface(
61
+ fn=predict_depth,
62
+ inputs=gr.Image(type="pil"),
63
+ outputs=[
64
+ gr.Image(type="pil", label="Original"),
65
+ gr.Image(type="pil", label="Depth Map")
66
+ ],
67
+ title="Depth Estimation API",
68
+ description="DepthAnything V2による深度推定"
69
+ )
 
 
 
 
 
 
 
 
 
 
 
 
 
 
70
 
 
71
  if __name__ == "__main__":
72
  demo.launch()