refactor(compressors): consolidate pipeline and improve mask handling

mini-nav/utils/image.py

@@ -9,7 +9,7 @@ def segment_image(
     image: Image.Image,
     min_area: int = 32 * 32,
     max_masks: int = 5,
-    points_per_batch=64,
+    points_per_batch: int = 64,
 ) -> list[dict[str, Any]]:
     """Segment image using SAM to extract object masks.
 
@@ -27,26 +27,104 @@ def segment_image(
             - predicted_iou: Model's confidence in the mask
             - stability_score: Stability score for the mask
     """
-    # Convert PIL Image to numpy array
-    image_np = np.array(image.convert("RGB"))
+    image_rgb = image.convert("RGB")
+    raw_output = mask_generator(image_rgb, points_per_batch=points_per_batch)
+    raw_masks = raw_output.get("masks", raw_output)
 
-    # Generate masks
-    masks = mask_generator(image_np, points_per_batch=points_per_batch)["masks"]
+    normalized_masks: list[dict[str, Any]] = []
 
-    if not masks:
+    if isinstance(raw_masks, list):
+        if raw_masks and isinstance(raw_masks[0], dict):
+            normalized_masks = raw_masks
+        else:
+            for mask_like in raw_masks:
+                mask_dict = _to_mask_dict(mask_like)
+                if mask_dict is not None:
+                    normalized_masks.append(mask_dict)
+    else:
+        mask_array = _to_numpy_mask_array(raw_masks)
+        if mask_array is not None:
+            if mask_array.ndim == 2:
+                mask_array = np.expand_dims(mask_array, axis=0)
+            if mask_array.ndim == 3:
+                for single_mask in mask_array:
+                    mask_dict = _to_mask_dict(single_mask)
+                    if mask_dict is not None:
+                        normalized_masks.append(mask_dict)
+
+    if not normalized_masks:
         return []
 
-    # Filter by minimum area
-    filtered_masks = [m for m in masks if m["area"] >= min_area]
+    filtered_masks = [m for m in normalized_masks if int(m["area"]) >= min_area]
 
     if not filtered_masks:
         return []
 
-    # Sort by area (largest first) and limit to max_masks
     sorted_masks = sorted(filtered_masks, key=lambda x: x["area"], reverse=True)
     return sorted_masks[:max_masks]
 
 
+def _to_numpy_mask_array(mask_like: Any) -> np.ndarray | None:
+    if mask_like is None:
+        return None
+    if isinstance(mask_like, np.ndarray):
+        return mask_like
+
+    try:
+        import torch
+
+        if isinstance(mask_like, torch.Tensor):
+            return mask_like.detach().cpu().numpy()
+    except ImportError:
+        pass
+
+    return None
+
+
+def _to_mask_dict(mask_like: Any) -> dict[str, Any] | None:
+    if isinstance(mask_like, dict):
+        if "area" in mask_like and "bbox" in mask_like and "segment" in mask_like:
+            return mask_like
+
+        segment = mask_like.get("segment")
+        if segment is None and "mask" in mask_like:
+            segment = mask_like["mask"]
+        if segment is None:
+            return None
+
+        mask_array = _to_numpy_mask_array(segment)
+        if mask_array is None:
+            return None
+        return _build_mask_dict(mask_array)
+
+    mask_array = _to_numpy_mask_array(mask_like)
+    if mask_array is None:
+        return None
+    return _build_mask_dict(mask_array)
+
+
+def _build_mask_dict(mask_array: np.ndarray) -> dict[str, Any] | None:
+    if mask_array.ndim != 2:
+        return None
+    segment = mask_array.astype(bool)
+    area = int(segment.sum())
+    if area <= 0:
+        return None
+
+    ys, xs = np.where(segment)
+    min_y, max_y = int(ys.min()), int(ys.max())
+    min_x, max_x = int(xs.min()), int(xs.max())
+    bbox = [min_x, min_y, max_x - min_x + 1, max_y - min_y + 1]
+
+    return {
+        "segment": segment,
+        "area": area,
+        "bbox": bbox,
+        "predicted_iou": None,
+        "stability_score": None,
+    }
+
+
 def extract_masked_region(
     image: Image.Image,
     mask: np.ndarray,