Update app.py

app.py

@@ -11,43 +11,32 @@ depth_estimator = pipeline(task="depth-estimation", model="Intel/dpt-hybrid-mida
 # Function to process the image and return depth map
 def launch(input_image):
     try:
-        # Resize the 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":
-            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
+        # Run the image segmentation model
         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)
+        # Assuming output contains the segmentation mask or predicted depth map
+        predicted_output = out["predicted_depth"] if "predicted_depth" in out else out["segmentation_mask"]
+        
+        # Resize the output to match the input image size
+        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)
             mode="bicubic",
-            align_corners=False,
+            align_corners=False
         )
 
-        # Normalize the prediction
-        output = prediction.squeeze().numpy()
+        # Normalize the prediction for display
+        output = predicted_output_resized.squeeze().cpu().numpy()
         formatted = (output * 255 / np.max(output)).astype("uint8")
-        
-        # Convert the depth map to an image
-        depth = Image.fromarray(formatted)
-        
-        return depth
+
+        # Convert the depth map or segmentation mask to an image
+        output_image = Image.fromarray(formatted)
+
+        return output_image
 
     except Exception as e:
         print(f"Error processing the image: {str(e)}")