Update app.py

app.py

@@ -5,23 +5,25 @@ import torch
 import numpy as np
 from PIL import Image
 
-# Load the depth estimation model
+# Load the depth estimation model or segmentation model
 depth_estimator = pipeline(task="depth-estimation", model="Intel/dpt-hybrid-midas")
 
-# Function to process the image and return depth map
 def launch(input_image):
     try:
+        # Resize the input image to a fixed size (e.g., 640x480)
+        input_image = input_image.resize((640, 480))
+        
         # Ensure the input image is in RGB mode
         if input_image.mode != "RGB":
             input_image = input_image.convert("RGB")
 
-        # Run the image segmentation model
+        # Run the image segmentation model (or depth estimation)
         out = depth_estimator(input_image)
         
-        # Assuming output contains the segmentation mask or predicted depth map
+        # Assuming the output contains the predicted depth or segmentation mask
         predicted_output = out["predicted_depth"] if "predicted_depth" in out else out["segmentation_mask"]
         
-        # Resize the output to match the input image size
+        # Resize the output to match the input image size (H, W)
         predicted_output_resized = torch.nn.functional.interpolate(
             predicted_output.unsqueeze(0),  # Add batch dimension
             size=input_image.size[::-1],  # Match input image size (H, W)
@@ -42,15 +44,13 @@ def launch(input_image):
         print(f"Error processing the image: {str(e)}")
         return "An error occurred while processing the image."
 
-# Define the Gradio interface
+# Define the Gradio interface without 'image_size' argument
 iface = gr.Interface(
     fn=launch,
     inputs=gr.Image(type="pil"),
-    outputs=gr.Image(type="pil", image_size=(640, 480))  # Specify size for output
+    outputs=gr.Image(type="pil")  # Remove image_size argument
 )
 
-
-
 # Launch the interface
 iface.launch()