refactor(pipeline): support multi-object selection in frame processing

mini-nav/compressors/pipeline.py

@@ -12,7 +12,6 @@ from utils.image import crop_image_by_bbox, extract_masked_region
 from .filter import (
     MaskScoringConfig,
     compute_mask_features,
-    score_mask,
     should_reject_mask,
 )
 from .model_loader import (
@@ -46,16 +45,18 @@ def create_pipeline_from_config(config) -> "HashPipeline":
 
 @dataclass
 class FramePacket:
+    """Multi-object frame packet for detection and crop outputs."""
+
     image: Image.Image
     boxes_xyxy: list[list[float]] = field(default_factory=list)
     scores: list[float] = field(default_factory=list)
     labels: list[str] = field(default_factory=list)
     masks: list[dict[str, Any]] = field(default_factory=list)
-    selected_idx: int | None = None
+    selected_indices: list[int] = field(default_factory=list)
     dropped_indices: list[int] = field(default_factory=list)
     fallback_reason: str | None = None
-    filtered_image: Image.Image | None = None
+    filtered_images: list[Image.Image] = field(default_factory=list)
-    cropped_image: Image.Image | None = None
+    cropped_images: list[Image.Image] = field(default_factory=list)
 
 
 @dataclass
@@ -196,13 +197,13 @@ class HashPipeline:
 
         for packet in packets:
             if not packet.masks:
-                packet.selected_idx = None
+                packet.selected_indices = []
                 packet.dropped_indices = []
                 if packet.fallback_reason is None:
                     packet.fallback_reason = "no_mask"
                 continue
 
-            kept: list[tuple[float, int, int]] = []
+            selected_indices: list[int] = []
             dropped: list[int] = []
             for index, mask in enumerate(packet.masks):
                 features = compute_mask_features(
@@ -212,49 +213,28 @@ class HashPipeline:
                     dropped.append(index)
                     continue
 
-                mask_score = score_mask(
+                selected_indices.append(index)
-                    mask,
-                    image_shape=(packet.image.height, packet.image.width),
-                    config=config,
-                )
-                area = int(mask.get("area", 0))
-                kept.append((float(mask_score), area, index))
 
-            if kept:
+            packet.selected_indices = selected_indices
-                kept.sort(reverse=True)
+            packet.dropped_indices = dropped
-                packet.selected_idx = kept[0][2]
+            if not selected_indices and packet.fallback_reason is None:
-                packet.dropped_indices = dropped
+                packet.fallback_reason = "all_masks_rejected"
-                continue
-
-            fallback_index = max(
-                range(len(packet.masks)),
-                key=lambda idx: int(packet.masks[idx].get("area", 0)),
-            )
-            packet.selected_idx = fallback_index
-            packet.dropped_indices = [
-                index for index in range(len(packet.masks)) if index != fallback_index
-            ]
-            if packet.fallback_reason is None:
-                packet.fallback_reason = "all_masks_rejected_fallback_area"
 
     def _render_filtered_images(self, packets: list[FramePacket]) -> None:
         for packet in packets:
-            if packet.selected_idx is None:
+            packet.filtered_images = []
-                packet.filtered_image = packet.image
-                continue
 
-            if packet.selected_idx >= len(packet.masks):
+            for selected_idx in packet.selected_indices:
-                packet.filtered_image = packet.image
+                if selected_idx >= len(packet.masks):
-                packet.fallback_reason = (
+                    packet.fallback_reason = (
-                    packet.fallback_reason or "selected_index_out_of_mask_range"
+                        packet.fallback_reason or "selected_index_out_of_mask_range"
+                    )
+                    continue
+
+                selected_mask = packet.masks[selected_idx]
+                packet.filtered_images.append(
+                    extract_masked_region(packet.image, selected_mask["segment"])
                 )
-                packet.selected_idx = None
-                continue
-
-            selected_mask = packet.masks[packet.selected_idx]
-            packet.filtered_image = extract_masked_region(
-                packet.image, selected_mask["segment"]
-            )
 
     def _render_cropped_images(
         self,
@@ -262,33 +242,16 @@ class HashPipeline:
         detections_per_image: list[list[DetectionResult]],
     ) -> None:
         for index, packet in enumerate(packets):
-            base_image = (
-                packet.filtered_image if packet.filtered_image else packet.image
-            )
             detections = (
                 detections_per_image[index] if index < len(detections_per_image) else []
             )
+            packet.cropped_images = []
 
-            if packet.selected_idx is None:
+            for local_idx, sel_idx in enumerate(packet.selected_indices):
-                packet.cropped_image = base_image
+                filtered_img = packet.filtered_images[local_idx]
-                if packet.fallback_reason is None:
+                selected_detection = detections[sel_idx]
-                    packet.fallback_reason = "no_selected_candidate"
+                cropped = crop_image_by_bbox(filtered_img, selected_detection["bbox"])
-                continue
+                packet.cropped_images.append(cropped)
-
-            if packet.selected_idx >= len(detections):
-                packet.cropped_image = base_image
-                packet.fallback_reason = (
-                    packet.fallback_reason or "selected_index_out_of_detection_range"
-                )
-                continue
-
-            selected_detection = detections[packet.selected_idx]
-            cropped = crop_image_by_bbox(base_image, selected_detection["bbox"])
-            packet.cropped_image = cropped
-            if cropped.size == base_image.size:
-                packet.fallback_reason = (
-                    packet.fallback_reason or "invalid_or_full_bbox"
-                )
 
     def _build_debug_meta(
         self,
@@ -298,7 +261,8 @@ class HashPipeline:
         debug_meta: list[dict[str, Any]] = []
         for packet in packets:
             item: dict[str, Any] = {
-                "selected_idx": packet.selected_idx,
+                "selected_indices": list(packet.selected_indices),
+                "num_selected": len(packet.selected_indices),
                 "dropped_indices": list(packet.dropped_indices),
                 "fallback_reason": packet.fallback_reason,
                 "num_boxes": len(packet.boxes_xyxy),
@@ -387,10 +351,20 @@ class HashPipeline:
         self._render_cropped_images(packets, detections_per_image)
 
         cropped_images = [
-            packet.cropped_image if packet.cropped_image is not None else packet.image
+            cropped_image
             for packet in packets
+            for cropped_image in packet.cropped_images
         ]
 
+        if not cropped_images:
+            return PipelineBatchOutput(
+                hash_bits=torch.empty(
+                    (0, self.hash_bits), dtype=torch.int32, device=self.device
+                ),
+                cropped_images=[],
+                debug_meta=self._build_debug_meta(packets, return_debug_details),
+            )
+
         all_bits: list[torch.Tensor] = []
         for index in range(0, len(cropped_images), batch_size):
             sub_batch = cropped_images[index : index + batch_size]