refactor(compressors): migrate pipeline to OWLv2-based detection with text labels

- Replace bbox-prompted segmentation with OWLv2 text-guided object detection
- Refactor HashPipeline from nn.Module to plain class with modular stage methods
- Add detect_batch, segment_batch, filter_batch for explicit pipeline stages
- Rename forward to forward_batch with text_labels API instead of bboxes
- Add mask_scoring_config, score_threshold, postprocess_threshold configuration
- Update model_loader to expose Dinov2Model type annotation

mini-nav/compressors/__init__.py

@@ -6,7 +6,7 @@ from .common import (
     hash_to_bits,
 )
 from .hash_compressor import HashCompressor, HashLoss, VideoPositiveMask
-from .object_score import (
+from .filter import (
     MaskFeatures,
     MaskScoringConfig,
     compute_mask_features,

mini-nav/compressors/model_loader.py

@@ -6,6 +6,7 @@ import torch
 from transformers import (
     AutoImageProcessor,
     AutoModel,
+    Dinov2Model,
     Owlv2ForObjectDetection,
     Owlv2Processor,
     Sam2Model,
@@ -34,7 +35,7 @@ def load_sam_model(
 
 def load_dino_model(
     model_name: str = "facebook/dinov2-large",
-) -> tuple[AutoImageProcessor, AutoModel]:
+) -> tuple[AutoImageProcessor, Dinov2Model]:
     device = get_device()
 
     processor = AutoImageProcessor.from_pretrained(model_name)

mini-nav/compressors/pipeline.py

@@ -1,25 +1,26 @@
-"""SAM + DINO + Hash compression pipeline."""
+"""OWLv2 + SAM + DINO + Hash compression pipeline."""
 
 from typing import Optional, Sequence
 
 import torch
-import torch.nn as nn
 import torch.nn.functional as F
 from PIL import Image
 from utils import get_device
 
-from .filter import select_best_mask
+from .filter import MaskScoringConfig, select_best_mask
 from .model_loader import (
     get_dino_dim,
     load_dino_model,
     load_hash_compressor,
+    load_owlv2_model,
     load_sam_model,
 )
 from .proposal import (
+    detect_objects_batch,
     extract_masked_region,
-    generate_proposals,
     generate_proposals_batch,
 )
+from .proposal.core import DetectionResult
 
 
 def create_pipeline_from_config(config) -> "HashPipeline":
@@ -32,50 +33,64 @@ def create_pipeline_from_config(config) -> "HashPipeline":
         Initialized HashPipeline.
     """
     return HashPipeline(
+        owlv2_model=getattr(
+            config.model, "owlv2_model", "google/owlv2-base-patch16-ensemble"
+        ),
         dino_model=config.model.dino_model,
         sam_model=config.model.sam_model,
-        sam_min_mask_area=config.model.sam_min_mask_area,
-        sam_max_masks=config.model.sam_max_masks,
         hash_bits=config.model.compression_dim,
         compressor_path=config.model.compressor_path,
+        mask_scoring_config=getattr(config.model, "mask_scoring_config", None),
+        score_threshold=getattr(config.model, "score_threshold", 0.25),
+        postprocess_threshold=getattr(config.model, "postprocess_threshold", 0.1),
     )
 
 
-class HashPipeline(nn.Module):
+class HashPipeline:
-    """Pipeline for SAM segmentation + DINO features + Hash compression.
+    """Pipeline for OWLv2 detection + SAM segmentation + DINO features + Hash compression.
 
     Pipeline flow:
-        PIL Image -> SAM (largest object mask) -> DINO (features) -> Hash (binary codes)
+        Images + Text Labels -> OWLv2 (detections) -> SAM (masks) -> Filter (best mask) ->
+        DINO (features) -> Hash (binary codes)
 
     Example:
         pipeline = HashPipeline(dino_model="facebook/dinov2-large", hash_bits=512)
-        image = Image.open("path/to/image.jpg")
+        images = [Image.open("path/to/image.jpg")]
-        hash_bits = pipeline(image)  # Returns [1, 512] binary bits
+        text_labels = ["object"]
+        hash_bits = pipeline.forward_batch(images, text_labels)  # Returns [N, 512]
     """
 
     def __init__(
         self,
+        owlv2_model: str = "google/owlv2-base-patch16-ensemble",
         dino_model: str = "facebook/dinov2-large",
         sam_model: str = "facebook/sam2.1-hiera-large",
-        sam_min_mask_area: int = 100,
-        sam_max_masks: int = 10,
         hash_bits: int = 512,
         compressor_path: Optional[str] = None,
+        mask_scoring_config: Optional["MaskScoringConfig"] = None,
+        score_threshold: float = 0.25,
+        postprocess_threshold: float = 0.1,
     ):
         super().__init__()
 
         # Device for model placement.
         self.device = get_device()
 
-        # SAM2 filter settings.
+        # OWLv2 detection settings.
-        self.sam_min_mask_area = sam_min_mask_area
+        self.owlv2_processor, self.owlv2_model = load_owlv2_model(
-        self.sam_max_masks = sam_max_masks
+            model_name=owlv2_model
+        )
+        self.score_threshold = score_threshold
+        self.postprocess_threshold = postprocess_threshold
 
-        # Load models.
+        # Mask scoring config for filter step.
+        self.mask_scoring_config = mask_scoring_config
+
+        # SAM2 model for segmentation.
         self.sam_processor, self.sam_model = load_sam_model(model_name=sam_model)
-        self.dino_processor, self.dino = load_dino_model(model_name=dino_model)
 
-        # DINO feature dimension based on model size.
+        # DINO model for feature extraction.
+        self.dino_processor, self.dino = load_dino_model(model_name=dino_model)
         self.dino_dim = get_dino_dim(dino_model)
 
         # Hash compressor for binarizing DINO features.
@@ -90,82 +105,112 @@ class HashPipeline(nn.Module):
         """Number of bits in the hash code."""
         return self.hash_compressor.hash_bits
 
-    def _segment_with_sam(
+    def detect_batch(
-        self, image: Image.Image, bboxes: list[list[float]]
+        self,
-    ) -> Image.Image:
+        images: Sequence[Image.Image],
-        """Segment image with SAM and extract the best object mask.
+        text_labels: list[str],
+    ) -> list[list[DetectionResult]]:
+        """Detect objects in a batch of images using OWLv2.
 
         Args:
-            image: Input PIL Image.
+            images: Sequence of PIL Images.
-            bboxes: Bounding boxes from object detector as [[x1,y1,x2,y2], ...].
+            text_labels: Text labels for all images (same labels used for each image).
 
         Returns:
-            Masked image containing only the best object, or original if no masks.
+            List of lists of DetectionResult dicts, one inner list per image.
         """
-        masks = generate_proposals(
+        image_list = list(images)
-            self.sam_model,
+        if not image_list:
-            self.sam_processor,
+            return []
-            image,
+
-            bboxes,
+        text_labels_per_image = [text_labels] * len(image_list)
+        return detect_objects_batch(
+            self.owlv2_model,
+            self.owlv2_processor,
+            image_list,
+            text_labels_per_image,
+            score_threshold=self.score_threshold,
+            postprocess_threshold=self.postprocess_threshold,
         )
 
-        masks = _filter_masks(masks, self.sam_min_mask_area, self.sam_max_masks)
+    def segment_batch(
-
-        if not masks:
-            return image
-
-        best_mask = select_best_mask(masks, image_shape=(image.height, image.width))
-        if best_mask is None:
-            return image
-        return extract_masked_region(image, best_mask["segment"])
-
-    def _segment_with_sam_dataset(
         self,
         images: Sequence[Image.Image],
         bboxes_per_image: list[list[list[float]]],
-    ) -> list[Image.Image]:
+    ) -> list[list[dict]]:
+        """Segment objects in images using SAM2 with bounding box prompts.
+
+        Args:
+            images: Sequence of PIL Images.
+            bboxes_per_image: Bounding boxes per image as [[[x1,y1,x2,y2], ...], ...].
+
+        Returns:
+            List of lists of mask dictionaries, one inner list per image.
+        """
         image_list = list(images)
-        masks_dataset = generate_proposals_batch(
+        if not image_list:
+            return []
+
+        return generate_proposals_batch(
             self.sam_model,
             self.sam_processor,
             image_list,
             bboxes_per_image,
         )
-        masks_dataset = [
-            _filter_masks(masks, self.sam_min_mask_area, self.sam_max_masks)
-            for masks in masks_dataset
-        ]
-        selected_images: list[Image.Image] = []
-        for image, masks in zip(image_list, masks_dataset):
-            if not masks:
-                selected_images.append(image)
-                continue
 
-            best_mask = select_best_mask(masks, image_shape=(image.height, image.width))
+    def filter_batch(
-            if best_mask is None:
+        self,
-                selected_images.append(image)
+        images: Sequence[Image.Image],
-                continue
+        masks_per_image: list[list[dict]],
-            selected_images.append(extract_masked_region(image, best_mask["segment"]))
+    ) -> list[Image.Image]:
-
+        """Filter masks and extract best masked region for each image.
-        return selected_images
-
-    def _dino_forward(self, image: Image.Image) -> torch.Tensor:
-        """Extract DINO tokens from an image.
 
         Args:
-            image: Input PIL Image.
+            images: Sequence of PIL Images.
+            masks_per_image: Masks per image from segment_batch.
 
         Returns:
-            Last hidden state tokens of shape [1, N, dim].
+            List of PIL Images, one per input image (original if no valid masks).
         """
-        inputs = self.dino_processor(image, return_tensors="pt").to(self.device)
+        image_list = list(images)
+        if not image_list:
+            return []
 
-        with torch.no_grad():
+        filtered_images: list[Image.Image] = []
-            outputs = self.dino(**inputs)
+        for image, masks in zip(image_list, masks_per_image):
-            return outputs.last_hidden_state
+            if not masks:
+                filtered_images.append(image)
+                continue
 
-    def _dino_forward_batch(self, images: Sequence[Image.Image]) -> torch.Tensor:
+            best_mask = select_best_mask(
-        inputs = self.dino_processor(images=list(images), return_tensors="pt").to(
+                masks,
+                image_shape=(image.height, image.width),
+                config=self.mask_scoring_config,
+            )
+            if best_mask is None:
+                filtered_images.append(image)
+                continue
+
+            filtered_images.append(extract_masked_region(image, best_mask["segment"]))
+
+        return filtered_images
+
+    def extract_dino_batch(self, images: Sequence[Image.Image]) -> torch.Tensor:
+        """Extract DINO tokens from a batch of images.
+
+        Args:
+            images: Sequence of PIL Images.
+
+        Returns:
+            Last hidden state tokens of shape [B, N, dim].
+        """
+        image_list = list(images)
+        if not image_list:
+            return torch.empty(
+                (0, 1, self.dino_dim), dtype=torch.float32, device=self.device
+            )
+
+        inputs = self.dino_processor(images=image_list, return_tensors="pt").to(
             self.device
         )
 
@@ -173,113 +218,92 @@ class HashPipeline(nn.Module):
             outputs = self.dino(**inputs)
             return outputs.last_hidden_state
 
-    def forward(self, image: Image.Image, bboxes: list[list[float]]) -> torch.Tensor:
+    def compress_batch(self, tokens: torch.Tensor) -> torch.Tensor:
-        """Process a single image through the full pipeline.
+        """Compress DINO tokens to binary hash codes.
 
         Args:
-            image: Input PIL Image.
+            tokens: DINO tokens of shape [B, N, dim].
-            bboxes: Bounding boxes from object detector as [[x1,y1,x2,y2], ...].
 
         Returns:
-            Binary hash codes of shape [1, hash_bits] as int32.
+            Binary hash codes of shape [B, hash_bits] as int32.
         """
-        image = self._segment_with_sam(image, bboxes)
-        tokens = self._dino_forward(image)
         _, _, bits = self.hash_compressor(tokens)
         return bits
 
-    def forward_dataset(
+    def forward_batch(
         self,
         images: Sequence[Image.Image],
-        bboxes_per_image: list[list[list[float]]],
+        text_labels: list[str],
         batch_size: int = 32,
-        apply_sam: bool = True,
     ) -> torch.Tensor:
+        """Process a batch of images through the full pipeline.
+
+        Args:
+            images: Sequence of PIL Images.
+            text_labels: Text labels for detection (same for all images).
+            batch_size: Batch size for DINO feature extraction.
+
+        Returns:
+            Binary hash codes of shape [N, hash_bits] as int32.
+        """
         if batch_size <= 0:
             raise ValueError("batch_size must be greater than 0")
 
         image_list = list(images)
-
+        if not image_list:
-        if len(image_list) == 0:
             return torch.empty(
                 (0, self.hash_bits), dtype=torch.int32, device=self.device
             )
 
-        if apply_sam:
+        detections = self.detect_batch(image_list, text_labels)
-            processed_images = self._segment_with_sam_dataset(
+        bboxes = [[d["bbox"] for d in dets] for dets in detections]
-                image_list, bboxes_per_image
+        masks = self.segment_batch(image_list, bboxes)
-            )
+        processed = self.filter_batch(image_list, masks)
-        else:
-            processed_images = image_list
 
-        batch_bits: list[torch.Tensor] = []
+        all_bits: list[torch.Tensor] = []
-        for i in range(0, len(processed_images), batch_size):
+        for i in range(0, len(processed), batch_size):
-            batch_images = processed_images[i : i + batch_size]
+            sub_batch = processed[i : i + batch_size]
-            tokens = self._dino_forward_batch(batch_images)
+            tokens = self.extract_dino_batch(sub_batch)
-            _, _, bits = self.hash_compressor(tokens)
+            bits = self.compress_batch(tokens)
-            batch_bits.append(bits)
+            all_bits.append(bits)
 
-        return torch.cat(batch_bits, dim=0)
+        return torch.cat(all_bits, dim=0)
-
-    def extract_features(
-        self, image: Image.Image, bboxes: list[list[float]]
-    ) -> torch.Tensor:
-        """Extract normalized DINO features from an image.
-
-        Args:
-            image: Input PIL Image.
-            bboxes: Bounding boxes from object detector as [[x1,y1,x2,y2], ...].
-
-        Returns:
-            Normalized DINO features of shape [1, dino_dim].
-        """
-        image = self._segment_with_sam(image, bboxes)
-        tokens = self._dino_forward(image)
-        features = tokens.mean(dim=1)
-        return F.normalize(features, dim=-1)
 
     def extract_features_dataset(
         self,
         images: Sequence[Image.Image],
-        bboxes_per_image: list[list[list[float]]],
+        text_labels: list[str],
         batch_size: int = 32,
-        apply_sam: bool = True,
     ) -> torch.Tensor:
+        """Extract normalized DINO features from a batch of images.
+
+        Args:
+            images: Sequence of PIL Images.
+            text_labels: Text labels for detection (same for all images).
+            batch_size: Batch size for DINO feature extraction.
+
+        Returns:
+            Normalized DINO features of shape [N, dino_dim].
+        """
         if batch_size <= 0:
             raise ValueError("batch_size must be greater than 0")
 
         image_list = list(images)
-
+        if not image_list:
-        if len(image_list) == 0:
             return torch.empty(
                 (0, self.dino_dim), dtype=torch.float32, device=self.device
             )
 
-        if apply_sam:
+        detections = self.detect_batch(image_list, text_labels)
-            processed_images = self._segment_with_sam_dataset(
+        bboxes = [[d["bbox"] for d in dets] for dets in detections]
-                image_list, bboxes_per_image
+        masks = self.segment_batch(image_list, bboxes)
-            )
+        processed = self.filter_batch(image_list, masks)
-        else:
-            processed_images = image_list
 
         all_features: list[torch.Tensor] = []
-        for i in range(0, len(processed_images), batch_size):
+        for i in range(0, len(processed), batch_size):
-            batch_images = processed_images[i : i + batch_size]
+            sub_batch = processed[i : i + batch_size]
-            tokens = self._dino_forward_batch(batch_images)
+            tokens = self.extract_dino_batch(sub_batch)
             features = tokens.mean(dim=1)
             all_features.append(F.normalize(features, dim=-1))
 
         return torch.cat(all_features, dim=0)
-
-
-def _filter_masks(
-    masks: list[dict],
-    min_area: int,
-    max_masks: int,
-) -> list[dict]:
-    """Filter masks by area and keep top-N largest."""
-    filtered = [m for m in masks if int(m["area"]) >= min_area]
-    if not filtered:
-        return []
-    sorted_masks = sorted(filtered, key=lambda m: m["area"], reverse=True)
-    return sorted_masks[:max_masks]