Update app.py

app.py

@@ -2,12 +2,12 @@ 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
-
+import cv2
+from huggingface_hub import hf_hub_download
 from depth_anything_v2.dpt import DepthAnythingV2
 
-# Load model as before
+# Model loading (as before)
 DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
 model_configs = {
     'vitl': {'encoder': 'vitl', 'features': 256, 'out_channels': [256, 512, 1024, 1024]},
@@ -23,21 +23,28 @@ state_dict = torch.load(model_path, map_location="cpu")
 model.load_state_dict(state_dict)
 model = model.to(DEVICE).eval()
 
-# Use a matplotlib colormap
 CMAP = plt.get_cmap('Spectral_r')
 
 def infer(image: np.ndarray):
-    # 1. Run the model (BGR to RGB if needed)
+    # Run depth model
     with torch.no_grad():
         depth = model.infer_image(image[:, :, ::-1])
-    # 2. Grayscale map (normalize to 0..255)
+
+    # 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 map
     colored = (CMAP(depth_norm)[:, :, :3] * 255).astype(np.uint8)
     color = Image.fromarray(colored)
-    return gray, color
+
+    # Edge map using Canny on the original image (convert to grayscale first)
+    image_gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) if image.shape[2] == 3 else image
+    edges = cv2.Canny(image_gray, 100, 200)  # threshold1/2 can be tuned
+    edge_img = Image.fromarray(edges)
+
+    return gray, color, edge_img
 
 iface = gr.Interface(
     fn=infer,
@@ -45,9 +52,10 @@ iface = gr.Interface(
     outputs=[
         gr.Image(label="Grayscale Depth"),
         gr.Image(label="Colored Depth"),
+        gr.Image(label="Canny Edge Map"),
     ],
-    title="Depth Anything V2 (Minimal, with Colored Output)",
-    description="Upload an image, get depth as grayscale and colored."
+    title="Depth Anything V2 (with Colored Output + Canny Edges)",
+    description="Upload an image to get depth (gray, color), plus Canny edge map."
 )
 
 iface.launch()