Manage image resizing

handler.py

@@ -64,50 +64,39 @@ class EndpointHandler():
         
         # 4. Process and select the best mask
         try:
-            # Use the processor's post_process_masks method correctly
-            post_processed_masks = self.processor.post_process_masks(
-                outputs.pred_masks,
-                inputs["original_sizes"],
-                inputs["reshaped_input_sizes"]
-            )
+            # Get original image dimensions
+            original_height, original_width = img.size[1], img.size[0]
+            
+            # Get predicted masks and scores
+            predicted_masks = outputs.pred_masks.cpu()
+            iou_scores = outputs.iou_scores.cpu()[0]
             
-            # post_processed_masks is a list with one element (for batch size 1)
-            masks = post_processed_masks[0]  # Shape: (num_masks, H, W)
-            iou_scores = outputs.iou_scores[0]  # Shape: (num_masks,)
+            # Handle different tensor dimensions
+            if predicted_masks.ndim == 5:
+                predicted_masks = predicted_masks.squeeze(1)
             
-            print("Number of masks generated: {}".format(masks.shape[0]))
-            print("IoU scores: {}".format(iou_scores.tolist()))
+            # Resize masks to standard size first
+            predicted_masks = torch.nn.functional.interpolate(
+                predicted_masks,
+                size=(1024, 1024),
+                mode='bilinear',
+                align_corners=False
+            )
             
-            # 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)
-                
-                # Safely convert to 2D by squeezing all singleton dimensions
-                best_mask = best_mask.squeeze()
-                
-                # If still not 2D, take the last 2 dimensions
-                if best_mask.ndim > 2:
-                    # Take the last 2 dimensions which should be height and width
-                    best_mask = best_mask.view(-1, best_mask.shape[-1])
-                elif best_mask.ndim == 1:
-                    # If somehow we got 1D, try to reshape to square
-                    size = int(best_mask.shape[0] ** 0.5)
-                    if size * size == best_mask.shape[0]:
-                        best_mask = best_mask.view(size, size)
-                    else:
-                        raise ValueError("Cannot reshape 1D mask to 2D")
-                
-                print("Final mask shape: {}".format(best_mask.shape))
-            else:
-                raise ValueError("No masks were generated")
+            # Select the best mask
+            best_mask_idx = torch.argmax(iou_scores)
+            best_mask = predicted_masks[0, best_mask_idx, :, :]
             
-            # Convert to numpy and create binary mask
-            mask_binary = (best_mask > 0.0).cpu().numpy().astype(np.uint8) * 255
+            # Resize to original image dimensions
+            final_mask = torch.nn.functional.interpolate(
+                best_mask.unsqueeze(0).unsqueeze(0),
+                size=(original_height, original_width),
+                mode='bilinear',
+                align_corners=False
+            ).squeeze()
             
-            print("Final mask_binary shape: {}".format(mask_binary.shape))
+            # Convert to binary mask
+            mask_binary = (final_mask > 0.0).numpy().astype(np.uint8) * 255
         
         except Exception as e:
             print("Error processing masks: {}".format(e))