Update app.py

app.py

@@ -11,6 +11,7 @@ from huggingface_hub import hf_hub_download
 from safetensors.torch import load_file
 from gradio_imageslider import ImageSlider
 import spaces
+import tempfile
 
 # Helper function to load model from Hugging Face
 def load_model_by_name(arch_name, checkpoint_path, device):
@@ -31,7 +32,7 @@ def load_model_by_name(arch_name, checkpoint_path, device):
 # Image processing function
 def process_image(image, model, device):
     if model is None:
-        return None
+        return None, None, None, None
     
     # Preprocess the image
     image_np = np.array(image)[..., ::-1] / 255
@@ -45,32 +46,40 @@ def process_image(image, model, device):
     image_tensor = transform({'image': image_np})['image']
     image_tensor = torch.from_numpy(image_tensor).unsqueeze(0).to(device)
     
-    with torch.no_grad():  # Disable autograd since we don't need gradients on CPU
+    with torch.no_grad():
         pred_disp, _ = model(image_tensor)
+    torch.cuda.empty_cache()
 
-    # Ensure the depth map is in the correct shape before colorization
-    pred_disp_np = pred_disp.cpu().detach().numpy()[0, 0, :, :]  # Remove extra singleton dimensions
+    # Convert depth map to numpy
+    pred_disp_np = pred_disp.cpu().detach().numpy()[0, 0, :, :]
     
     # Normalize depth map
-    pred_disp = (pred_disp_np - pred_disp_np.min()) / (pred_disp_np.max() - pred_disp_np.min())
+    pred_disp_normalized = (pred_disp_np - pred_disp_np.min()) / (pred_disp_np.max() - pred_disp_np.min())
     
-    # Colorize depth map
+    # Colorized depth map
     cmap = "Spectral_r"
-    depth_colored = colorize_depth_maps(pred_disp[None, ..., None], 0, 1, cmap=cmap).squeeze()  # Ensure correct dimension
-    
-    # Convert to uint8 for image display
+    depth_colored = colorize_depth_maps(pred_disp_normalized[None, ..., None], 0, 1, cmap=cmap).squeeze()
     depth_colored = (depth_colored * 255).astype(np.uint8)
-    
-    # Convert to HWC format (height, width, channels)
     depth_colored_hwc = chw2hwc(depth_colored)
     
-    # Resize to match the original image dimensions (height, width)
+    # Gray depth map
+    depth_gray = (pred_disp_normalized * 255).astype(np.uint8)
+    depth_gray_hwc = np.stack([depth_gray] * 3, axis=-1)  # Convert to 3-channel grayscale
+    
+    # Save raw depth map as a temporary npy file
+    with tempfile.NamedTemporaryFile(delete=False, suffix=".npy") as temp_file:
+        np.save(temp_file.name, pred_disp_normalized)
+        depth_raw_path = temp_file.name
+    
+    # Resize outputs to match original image size
     h, w = image_np.shape[:2]
     depth_colored_hwc = cv2.resize(depth_colored_hwc, (w, h), cv2.INTER_LINEAR)
+    depth_gray_hwc = cv2.resize(depth_gray_hwc, (w, h), cv2.INTER_LINEAR)
     
-    # Convert to a PIL image
-    depth_image = Image.fromarray(depth_colored_hwc)
-    return image, depth_image
+    # Convert to PIL images
+    return image, Image.fromarray(depth_colored_hwc), Image.fromarray(depth_gray_hwc), depth_raw_path
+
+
 
 # Gradio interface function with GPU support
 @spaces.GPU
@@ -105,17 +114,20 @@ def gradio_interface(image):
     model = model.to(device)  # 确保模型在正确的设备上
     
     if model is None:
-        return None
+        return None, None, None, None
     
     # Process image and return output
-    depth_image = process_image(image, model, device)
-    return depth_image
+    image, depth_image, depth_gray, depth_raw = process_image(image, model, device)
+    return (image, depth_image), depth_gray, depth_raw
 
 # Create Gradio interface
 iface = gr.Interface(
     fn=gradio_interface,
     inputs=gr.Image(type="pil"),  # Only image input, no mode selection
-    outputs = ImageSlider(label="Depth slider", type="pil", slider_color="pink"), # Depth image out with a slider
+    outputs = [ImageSlider(label="Depth slider", type="pil", slider_color="pink"), # Depth image out with a slider
+        gr.Image(type="pil", label="Gray Depth"),
+        gr.File(label="Raw Depth (NumPy File)")
+    ],
     title="Depth Estimation Demo",
     description="Upload an image to see the depth estimation results. Our model is running on GPU for faster processing.",
     examples=["1.jpg", "2.jpg", "4.png", "5.jpg", "6.jpg"],