Manage image resizing

handler.py

@@ -53,8 +53,9 @@ class EndpointHandler():
         
         # 2. Prepare prompts and process the image
         height, width = img.size[1], img.size[0]
-        input_points = [[[width // 2, height // 2]]]
-        input_labels = [[1]]
+        # Use multiple points for better segmentation
+        input_points = [[[width // 4, height // 4], [width // 2, height // 2], [3 * width // 4, 3 * height // 4]]]
+        input_labels = [[1, 1, 1]]  # All positive points
         
         inputs = self.processor(img, input_points=input_points, input_labels=input_labels, return_tensors="pt").to(device)
         
@@ -65,7 +66,6 @@ class EndpointHandler():
         # 4. Process and select the best mask
         try:
             # Use the processor's post_process_masks method correctly
-            # This method expects the raw model outputs and the input metadata
             post_processed_masks = self.processor.post_process_masks(
                 outputs.pred_masks,
                 inputs["original_sizes"],
@@ -73,13 +73,20 @@ class EndpointHandler():
             )
             
             # post_processed_masks is a list with one element (for batch size 1)
-            # Each element has shape (num_masks, original_height, original_width)
             masks = post_processed_masks[0]  # Shape: (num_masks, H, W)
-            
-            # Get IoU scores and select the best mask
             iou_scores = outputs.iou_scores[0]  # Shape: (num_masks,)
-            best_mask_idx = torch.argmax(iou_scores)
-            best_mask = masks[best_mask_idx]  # Shape: (H, W)
+            
+            print("Number of masks generated: {}".format(masks.shape[0]))
+            print("IoU scores: {}".format(iou_scores.tolist()))
+            
+            # Ensure we have masks and select the best one safely
+            if masks.shape[0] > 0:
+                best_mask_idx = torch.argmax(iou_scores)
+                # Ensure the index is within bounds
+                best_mask_idx = min(best_mask_idx.item(), masks.shape[0] - 1)
+                best_mask = masks[best_mask_idx]  # Shape: (H, W)
+            else:
+                raise ValueError("No masks were generated")
             
             # Convert to numpy and create binary mask
             mask_binary = (best_mask > 0.0).cpu().numpy().astype(np.uint8) * 255