Vanyadoing commited on
Commit
b565be2
·
verified ·
1 Parent(s): 9e7ca7c

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +9 -6
app.py CHANGED
@@ -1,12 +1,9 @@
1
  import gradio as gr
2
  import torch
 
3
  from huggingface_hub import hf_hub_download
4
  from depth_anything_v2.dpt import DepthAnythingV2
5
 
6
- def dummy_infer(img):
7
- return 255 - img
8
-
9
- # --- LOAD THE MODEL, BUT DON'T USE IT ---
10
  DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
11
  model_configs = {
12
  'vitl': {'encoder': 'vitl', 'features': 256, 'out_channels': [256, 512, 1024, 1024]},
@@ -21,7 +18,13 @@ model_path = hf_hub_download(
21
  state_dict = torch.load(model_path, map_location="cpu")
22
  model.load_state_dict(state_dict)
23
  model = model.to(DEVICE).eval()
24
- # --- END MODEL LOADING ---
25
 
26
- iface = gr.Interface(fn=dummy_infer, inputs=gr.Image(type="numpy"), outputs=gr.Image())
 
 
 
 
 
 
 
27
  iface.launch()
 
1
  import gradio as gr
2
  import torch
3
+ import numpy as np
4
  from huggingface_hub import hf_hub_download
5
  from depth_anything_v2.dpt import DepthAnythingV2
6
 
 
 
 
 
7
  DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
8
  model_configs = {
9
  'vitl': {'encoder': 'vitl', 'features': 256, 'out_channels': [256, 512, 1024, 1024]},
 
18
  state_dict = torch.load(model_path, map_location="cpu")
19
  model.load_state_dict(state_dict)
20
  model = model.to(DEVICE).eval()
 
21
 
22
+ def infer(img):
23
+ with torch.no_grad():
24
+ depth = model.infer_image(img[:, :, ::-1]) # BGR to RGB if needed
25
+ # Normalize to 0-255 and convert to uint8
26
+ depth_norm = (depth - depth.min()) / (depth.max() - depth.min()) * 255.0
27
+ return depth_norm.astype(np.uint8)
28
+
29
+ iface = gr.Interface(fn=infer, inputs=gr.Image(type="numpy"), outputs=gr.Image())
30
  iface.launch()