Manage image resizing

handler.py

@@ -85,11 +85,23 @@ class EndpointHandler():
                 # 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)
+                
+                # Ensure the mask is 2D by squeezing any extra dimensions
+                while best_mask.ndim > 2:
+                    best_mask = best_mask.squeeze()
+                
+                print("Final mask shape: {}".format(best_mask.shape))
             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
+            
+            # Ensure mask_binary is 2D
+            if mask_binary.ndim > 2:
+                mask_binary = mask_binary.squeeze()
+            
+            print("Final mask_binary shape: {}".format(mask_binary.shape))
         
         except Exception as e:
             print("Error processing masks: {}".format(e))