distill-any-depth

app.py

@@ -3,9 +3,9 @@ import torch
 from PIL import Image
 import cv2
 import numpy as np
-from geobench.modeling.archs.dam.dam import DepthAnything
-from geobench.utils.image_util import colorize_depth_maps
-from geobench.midas.transforms import Resize, NormalizeImage, PrepareForNet
+from distillanydepth.modeling.archs.dam.dam import DepthAnything
+from distillanydepth.utils.image_util import colorize_depth_maps
+from distillanydepth.midas.transforms import Resize, NormalizeImage, PrepareForNet
 from torchvision.transforms import Compose
 import os
 
@@ -21,7 +21,7 @@ def load_model_by_name(arch_name, checkpoint_path, device):
     return model
 
 # Image processing function (same as your original code, modified for Gradio)
-def process_image(image, model, device, mode='rel_depth'):
+def process_image(image, model, device):
     # Preprocess the image
     image_np = np.array(image)[..., ::-1] / 255
     transform = Compose([
@@ -39,7 +39,7 @@ def process_image(image, model, device, mode='rel_depth'):
         pred_disp = (pred_disp_np - pred_disp_np.min()) / (pred_disp_np.max() - pred_disp_np.min())
 
     # Colorize depth map
-    cmap = "Spectral_r" if mode != 'metric' else 'Spectral_r'
+    cmap = "Spectral_r"  # Default colormap for relative depth
     depth_colored = colorize_depth_maps(pred_disp[None, ...], 0, 1, cmap=cmap).squeeze()
     depth_colored = (depth_colored * 255).astype(np.uint8)
     
@@ -47,22 +47,22 @@ def process_image(image, model, device, mode='rel_depth'):
     return depth_image
 
 # Gradio interface function
-def gradio_interface(image, mode='rel_depth'):
+def gradio_interface(image):
     # Set device to CPU explicitly
     device = torch.device("cpu")  # Force using CPU
     model = load_model_by_name("depthanything", "your_checkpoint_path_here", device)
 
     # Process image and return output
-    return process_image(image, model, device, mode)
+    return process_image(image, model, device)
 
 # Create Gradio interface
 iface = gr.Interface(
     fn=gradio_interface,
-    inputs=[gr.Image(type="pil"), gr.Dropdown(choices=['rel_depth', 'metric_depth', 'disparity'], label="Mode")],
+    inputs=gr.Image(type="pil"),  # Only image input, no mode selection
     outputs=gr.Image(type="pil"),
     title="Depth Estimation Demo",
     description="Upload an image to see the depth estimation results."
 )
 
 # Launch the Gradio interface
-iface.launch()
+iface.launch()