Manage image resizing

handler.py

@@ -86,9 +86,20 @@ class EndpointHandler():
                 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()
+                # 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:
@@ -97,10 +108,6 @@ class EndpointHandler():
             # 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: