Update app.py

app.py

@@ -5,6 +5,7 @@ import numpy as np
 import matplotlib.pyplot as plt
 import subprocess
 import spaces
+import torch
 
 # Run the script to get pretrained models
 subprocess.run(["bash", "get_pretrained_models.sh"])
@@ -15,39 +16,52 @@ model.eval()
 
 @spaces.GPU(duration=120)
 def predict_depth(input_image):
-    # Preprocess the image
-    result = depth_pro.load_rgb(input_image)  # Removed .name
-    image = result[0]
-    f_px = result[-1]  # Assuming f_px is the last item in the returned tuple
-    image = transform(image)
-
-    # Run inference
-    prediction = model.infer(image, f_px=f_px)
-    depth = prediction["depth"]  # Depth in [m]
-    focallength_px = prediction["focallength_px"]  # Focal length in pixels
-
-    # Normalize depth for visualization
-    depth_normalized = (depth - np.min(depth)) / (np.max(depth) - np.min(depth))
-    
-    # Create a color map
-    plt.figure(figsize=(10, 10))
-    plt.imshow(depth_normalized, cmap='viridis')
-    plt.colorbar(label='Depth')
-    plt.title('Predicted Depth Map')
-    plt.axis('off')
-    
-    # Save the plot to a file
-    output_path = "depth_map.png"
-    plt.savefig(output_path)
-    plt.close()
-
-    return output_path, f"Focal length: {focallength_px:.2f} pixels"
+    try:
+        # Preprocess the image
+        result = depth_pro.load_rgb(input_image)
+        image = result[0]
+        f_px = result[-1]  # Assuming f_px is the last item in the returned tuple
+        image = transform(image)
+
+        # Run inference
+        prediction = model.infer(image, f_px=f_px)
+        depth = prediction["depth"]  # Depth in [m]
+        focallength_px = prediction["focallength_px"]  # Focal length in pixels
+
+        # Convert depth to numpy array if it's a torch tensor
+        if isinstance(depth, torch.Tensor):
+            depth = depth.cpu().numpy()
+
+        # Ensure depth is a 2D numpy array
+        if depth.ndim != 2:
+            depth = depth.squeeze()
+
+        # Normalize depth for visualization
+        depth_min = np.min(depth)
+        depth_max = np.max(depth)
+        depth_normalized = (depth - depth_min) / (depth_max - depth_min)
+        
+        # Create a color map
+        plt.figure(figsize=(10, 10))
+        plt.imshow(depth_normalized, cmap='viridis')
+        plt.colorbar(label='Depth')
+        plt.title('Predicted Depth Map')
+        plt.axis('off')
+        
+        # Save the plot to a file
+        output_path = "depth_map.png"
+        plt.savefig(output_path)
+        plt.close()
+
+        return output_path, f"Focal length: {focallength_px:.2f} pixels"
+    except Exception as e:
+        return None, f"An error occurred: {str(e)}"
 
 # Create Gradio interface
 iface = gr.Interface(
     fn=predict_depth,
     inputs=gr.Image(type="filepath"),
-    outputs=[gr.Image(type="filepath", label="Depth Map"), gr.Textbox(label="Focal Length")],
+    outputs=[gr.Image(type="filepath", label="Depth Map"), gr.Textbox(label="Focal Length or Error Message")],
     title="Depth Prediction Demo",
     description="Upload an image to predict its depth map and focal length."
 )