feat(verification): add batch segmentation and image saving

mini-nav/utils/image.py

@@ -29,10 +29,79 @@ def segment_image(
     """
     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)
+    return _normalize_and_filter_masks(
+        raw_output, min_area=min_area, max_masks=max_masks
+    )
+
+
+def segment_image_dataset(
+    mask_generator: Any,
+    images: Sequence[Image.Image],
+    min_area: int = 32 * 32,
+    max_masks: int = 5,
+    points_per_batch: int = 64,
+) -> list[list[dict[str, Any]]]:
+    image_list = list(images)
+    if not image_list:
+        return []
+
+    image_rgb_list = [image.convert("RGB") for image in image_list]
+    try:
+        raw_batch_output = mask_generator(
+            image_rgb_list,
+            points_per_batch=points_per_batch,
+        )
+        batch_items = _split_batch_output(
+            raw_batch_output, expected_size=len(image_list)
+        )
+        if batch_items is not None:
+            return [
+                _normalize_and_filter_masks(
+                    batch_item,
+                    min_area=min_area,
+                    max_masks=max_masks,
+                )
+                for batch_item in batch_items
+            ]
+    except TypeError:
+        pass
+
+    return [
+        _normalize_and_filter_masks(
+            mask_generator(image_rgb, points_per_batch=points_per_batch),
+            min_area=min_area,
+            max_masks=max_masks,
+        )
+        for image_rgb in image_rgb_list
+    ]
+
+
+def _split_batch_output(raw_output: Any, expected_size: int) -> list[Any] | None:
+    if isinstance(raw_output, list):
+        if len(raw_output) == expected_size:
+            return raw_output
+        return None
+
+    if isinstance(raw_output, dict):
+        raw_masks = raw_output.get("masks", raw_output)
+        if isinstance(raw_masks, list) and len(raw_masks) == expected_size:
+            return raw_masks
+
+    return None
+
+
+def _normalize_and_filter_masks(
+    raw_output: Any,
+    min_area: int,
+    max_masks: int,
+) -> list[dict[str, Any]]:
+    raw_masks = (
+        raw_output.get("masks", raw_output)
+        if isinstance(raw_output, dict)
+        else raw_output
+    )
 
     normalized_masks: list[dict[str, Any]] = []
-
     if isinstance(raw_masks, list):
         if raw_masks and isinstance(raw_masks[0], dict):
             normalized_masks = raw_masks
@@ -55,35 +124,16 @@ def segment_image(
     if not normalized_masks:
         return []
 
-    filtered_masks = [m for m in normalized_masks if int(m["area"]) >= min_area]
-
+    filtered_masks = [
+        mask for mask in normalized_masks if int(mask["area"]) >= min_area
+    ]
     if not filtered_masks:
         return []
 
-    sorted_masks = sorted(filtered_masks, key=lambda x: x["area"], reverse=True)
+    sorted_masks = sorted(filtered_masks, key=lambda mask: mask["area"], reverse=True)
     return sorted_masks[:max_masks]
 
 
-def segment_image_dataset(
-    mask_generator: Any,
-    images: Sequence[Image.Image],
-    min_area: int = 32 * 32,
-    max_masks: int = 5,
-    points_per_batch: int = 64,
-) -> list[list[dict[str, Any]]]:
-    image_list = list(images)
-    return [
-        segment_image(
-            mask_generator,
-            image,
-            min_area=min_area,
-            max_masks=max_masks,
-            points_per_batch=points_per_batch,
-        )
-        for image in image_list
-    ]
-
-
 def _to_numpy_mask_array(mask_like: Any) -> np.ndarray | None:
     if mask_like is None:
         return None