Update app.py

app.py

@@ -10,38 +10,46 @@ depth_estimator = pipeline(task="depth-estimation", model="Intel/dpt-hybrid-mida
 
 # Function to process the image and return depth map
 def launch(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}")
-
-    # Run depth estimation
-    out = depth_estimator(input_image)
-
-    # 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)
-        mode="bicubic",
-        align_corners=False,
-    )
-
-    # 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
+    try:
+        # Ensure the input image is in RGB mode
+        if input_image.mode != "RGB":
+            print(f"Converting image from {input_image.mode} to RGB.")
+            input_image = input_image.convert("RGB")
+
+        # Print input image details for debugging
+        print(f"Received image with size: {input_image.size}")
+
+        # Run depth estimation
+        out = depth_estimator(input_image)
+        print(f"Model output: {out}")
+
+        # 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)
+            mode="bicubic",
+            align_corners=False,
+        )
+
+        # 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
+
+    except Exception as e:
+        print(f"Error processing the image: {str(e)}")
+        return "An error occurred while processing the image."
+
 
 # Define the Gradio interface
 iface = gr.Interface(