Update app.py

app.py

@@ -10,14 +10,26 @@ depth_estimator = pipeline(task="depth-estimation", model="Intel/dpt-hybrid-mida
 
 # Function to process the image and return depth map
 def launch(input_image):
-    out = depth_estimator(input_image)
+    # Ensure the input image is in RGB mode
+    if input_image.mode != "RGB":
+        input_image = input_image.convert("RGB")
+    
+    # Print input image details for debugging
+    print(f"Received image with size: {input_image.size}")
 
-    predicted_depth = out["predicted_depth"].view(1, 1, 480, 640)  # Assuming single image.
+    # Run depth estimation
+    out = depth_estimator(input_image)
 
-    # Resize the prediction to match the raw image size (H, W).
+    # Check if the model output contains 'predicted_depth'
+    if "predicted_depth" in out:
+        predicted_depth = out["predicted_depth"].view(1, 1, 480, 640)  # Assuming single image
+    else:
+        raise ValueError("Model output does not contain 'predicted_depth'.")
+    
+    # Resize the prediction to match the raw image size (H, W)
     prediction = torch.nn.functional.interpolate(
         predicted_depth,
-        size=input_image.size[::-1],  # Match raw image size (H, W).
+        size=input_image.size[::-1],  # Match raw image size (H, W)
         mode="bicubic",
         align_corners=False,
     )
@@ -25,15 +37,19 @@ def launch(input_image):
     # Normalize the prediction
     output = prediction.squeeze().numpy()
     formatted = (output * 255 / np.max(output)).astype("uint8")
+    
+    # Convert the depth map to an image
     depth = Image.fromarray(formatted)
+    
     return depth
 
 # Define the Gradio interface
 iface = gr.Interface(
-    fn=launch, 
-    inputs=gr.Image(type='pil'), 
-    outputs=gr.Image(type='pil')
+    fn=launch,
+    inputs=gr.Image(type="pil"),  # Ensure input is PIL image
+    outputs=gr.Image(type="pil")  # Output is also in PIL format
 )
 
 # Launch the interface
 iface.launch()
+