Update app.py

app.py

@@ -1,9 +1,13 @@
 import gradio as gr
 import torch
 import numpy as np
+from PIL import Image
 from huggingface_hub import hf_hub_download
+import matplotlib.pyplot as plt
+
 from depth_anything_v2.dpt import DepthAnythingV2
 
+# Load model as before
 DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
 model_configs = {
     'vitl': {'encoder': 'vitl', 'features': 256, 'out_channels': [256, 512, 1024, 1024]},
@@ -19,12 +23,31 @@ state_dict = torch.load(model_path, map_location="cpu")
 model.load_state_dict(state_dict)
 model = model.to(DEVICE).eval()
 
-def infer(img):
+# Use a matplotlib colormap
+CMAP = plt.get_cmap('Spectral_r')
+
+def infer(image: np.ndarray):
+    # 1. Run the model (BGR to RGB if needed)
     with torch.no_grad():
-        depth = model.infer_image(img[:, :, ::-1])  # BGR to RGB if needed
-        # Normalize to 0-255 and convert to uint8
-        depth_norm = (depth - depth.min()) / (depth.max() - depth.min()) * 255.0
-        return depth_norm.astype(np.uint8)
+        depth = model.infer_image(image[:, :, ::-1])
+    # 2. Grayscale map (normalize to 0..255)
+    depth_norm = (depth - depth.min()) / (depth.max() - depth.min()) * 255.0
+    depth_norm = depth_norm.astype(np.uint8)
+    gray = Image.fromarray(depth_norm)
+    # 3. Color map
+    colored = (CMAP(depth_norm)[:, :, :3] * 255).astype(np.uint8)
+    color = Image.fromarray(colored)
+    return gray, color
+
+iface = gr.Interface(
+    fn=infer,
+    inputs=gr.Image(type="numpy", label="Input Image"),
+    outputs=[
+        gr.Image(label="Grayscale Depth"),
+        gr.Image(label="Colored Depth"),
+    ],
+    title="Depth Anything V2 (Minimal, with Colored Output)",
+    description="Upload an image, get depth as grayscale and colored."
+)
 
-iface = gr.Interface(fn=infer, inputs=gr.Image(type="numpy"), outputs=gr.Image())
 iface.launch()