Fix mask output: use decoder output[1] and resize to 1024x1024

handler.py

@@ -54,12 +54,14 @@ class EndpointHandler:
                 'point_labels': labels.reshape(1, -1)
             })
 
-            masks = decoder_outputs[0]
+            # decoder_outputs[0] is IoU scores (1, 4)
+            # decoder_outputs[1] is masks (1, 4, 256, 256)
+            masks = decoder_outputs[1]
 
-            # 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
+            # Take first mask and resize to 1024x1024
+            mask = masks[0, 0]  # Shape: (256, 256)
+            mask = Image.fromarray(mask).resize((1024, 1024), Image.BILINEAR)
+            mask = np.array(mask)
             mask = (mask > 0.0).astype(np.uint8) * 255
 
             # Return result