feat(pipeline): add batch processing for scene graph construction

mini-nav/compressors/pipeline.py

@@ -1,6 +1,6 @@
 """SAM + DINO + Hash compression pipeline."""
 
-from typing import Optional
+from typing import Optional, Sequence
 
 import torch
 import torch.nn as nn
@@ -8,7 +8,7 @@ import torch.nn.functional as F
 from PIL import Image
 
 from utils import get_device
-from utils.image import extract_masked_region, segment_image
+from utils.image import extract_masked_region, segment_image, segment_image_dataset
 from utils.model import (
     get_dino_dim,
     load_dino_model,
@@ -105,6 +105,7 @@ class HashPipeline(nn.Module):
             image,
             min_area=self.sam_min_mask_area,
             max_masks=self.sam_max_masks,
+            points_per_batch=self.sam_points_per_batch,
         )
 
         if not masks:
@@ -112,6 +113,23 @@ class HashPipeline(nn.Module):
 
         return extract_masked_region(image, masks[0]["segment"])
 
+    def _segment_with_sam_dataset(
+        self,
+        images: Sequence[Image.Image],
+    ) -> list[Image.Image]:
+        image_list = list(images)
+        masks_dataset = segment_image_dataset(
+            self.mask_generator,
+            image_list,
+            min_area=self.sam_min_mask_area,
+            max_masks=self.sam_max_masks,
+            points_per_batch=self.sam_points_per_batch,
+        )
+        return [
+            extract_masked_region(image, masks[0]["segment"]) if masks else image
+            for image, masks in zip(image_list, masks_dataset)
+        ]
+
     def _dino_forward(self, image: Image.Image) -> torch.Tensor:
         """Extract DINO tokens from an image.
 
@@ -127,6 +145,15 @@ class HashPipeline(nn.Module):
             outputs = self.dino(**inputs)
             return outputs.last_hidden_state
 
+    def _dino_forward_batch(self, images: Sequence[Image.Image]) -> torch.Tensor:
+        inputs = self.processor(images=list(images), return_tensors="pt").to(
+            self.device
+        )
+
+        with torch.no_grad():
+            outputs = self.dino(**inputs)
+            return outputs.last_hidden_state
+
     def forward(self, image: Image.Image) -> torch.Tensor:
         """Process a single image through the full pipeline.
 
@@ -141,6 +168,36 @@ class HashPipeline(nn.Module):
         _, _, bits = self.hash_compressor(tokens)
         return bits
 
+    def forward_dataset(
+        self,
+        images: Sequence[Image.Image],
+        batch_size: int = 32,
+        apply_sam: bool = True,
+    ) -> torch.Tensor:
+        if batch_size <= 0:
+            raise ValueError("batch_size must be greater than 0")
+
+        image_list = list(images)
+
+        if len(image_list) == 0:
+            return torch.empty(
+                (0, self.hash_bits), dtype=torch.int32, device=self.device
+            )
+
+        if apply_sam:
+            processed_images = self._segment_with_sam_dataset(image_list)
+        else:
+            processed_images = image_list
+
+        batch_bits: list[torch.Tensor] = []
+        for i in range(0, len(processed_images), batch_size):
+            batch_images = processed_images[i : i + batch_size]
+            tokens = self._dino_forward_batch(batch_images)
+            _, _, bits = self.hash_compressor(tokens)
+            batch_bits.append(bits)
+
+        return torch.cat(batch_bits, dim=0)
+
     def extract_features(self, image: Image.Image) -> torch.Tensor:
         """Extract normalized DINO features from an image.
 
@@ -154,3 +211,33 @@ class HashPipeline(nn.Module):
         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],
+        batch_size: int = 32,
+        apply_sam: bool = True,
+    ) -> torch.Tensor:
+        if batch_size <= 0:
+            raise ValueError("batch_size must be greater than 0")
+
+        image_list = list(images)
+
+        if len(image_list) == 0:
+            return torch.empty(
+                (0, self.dino_dim), dtype=torch.float32, device=self.device
+            )
+
+        if apply_sam:
+            processed_images = self._segment_with_sam_dataset(image_list)
+        else:
+            processed_images = image_list
+
+        all_features: list[torch.Tensor] = []
+        for i in range(0, len(processed_images), batch_size):
+            batch_images = processed_images[i : i + batch_size]
+            tokens = self._dino_forward_batch(batch_images)
+            features = tokens.mean(dim=1)
+            all_features.append(F.normalize(features, dim=-1))
+
+        return torch.cat(all_features, dim=0)