Fix mask dimension handling with squeeze

handler.py

@@ -48,14 +48,19 @@ class EndpointHandler:
             labels = np.array(params.get("point_labels", [1]), dtype=np.float32)
 
             # Decode
-            masks = self.decoder.run(None, {
+            decoder_outputs = self.decoder.run(None, {
                 'image_embeddings': embeddings,
                 'point_coords': coords.reshape(1, -1, 2),
                 'point_labels': labels.reshape(1, -1)
-            })[0]
+            })
 
-            # Postprocess
-            mask = (masks[0, 0] > 0.0).astype(np.uint8) * 255
+            masks = decoder_outputs[0]
+
+            # Postprocess - squeeze to get 2D mask
+            mask = masks.squeeze()  # Remove all dimensions of size 1
+            if len(mask.shape) > 2:
+                mask = mask[0]  # Take first mask if multiple
+            mask = (mask > 0.0).astype(np.uint8) * 255
 
             # Return result
             result = {"mask_shape": list(mask.shape), "has_object": bool(mask.max() > 0)}