SikongJueluo/Mini-Nav
1
0
Fork 0
mirror of https://github.com/SikongJueluo/Mini-Nav.git synced 2026-07-13 04:25:32 +08:00
Code Issues Packages Projects Releases Wiki Activity

refactor(compressors): switch SAM from automatic mask generation to bbox-prompted segmentation

Browse Source 
- Replace SAM2AutomaticMaskGenerator pipeline with Sam2Processor+Sam2Model
- Freeze SAM model parameters at load time, removing torch.no_grad() at call sites
- Rewrite proposal/core.py to use bbox prompts instead of automatic point sampling
- Add bboxes parameter to all HashPipeline public methods (forward, forward_dataset, extract_features, extract_features_dataset)
- Extract mask filtering logic (_filter_masks) from proposal into pipeline
- Rename object_score/ to filter/
- Add load_owlv2_model to model_loader
- Rename notebooks/test.py to habitat_sim_setup.py
This commit is contained in:
SikongJueluo
2026-04-02 16:42:59 +08:00
parent eaf02cc97a
commit 42acb3ee1b
9 changed files with 161 additions and 167 deletions
Download Patch File Download Diff File Expand all files Collapse all files

95
mini-nav/compressors/pipeline.py
View File

@@ -6,20 +6,20 @@ import torch
import torch.nn as nn
import torch.nn.functional as F
from PIL import Image
from .object_score import select_best_mask
from .proposal import (
extract_masked_region,
generate_proposals,
generate_proposals_batch,
)
from utils import get_device
from .filter import select_best_mask
from .model_loader import (
get_dino_dim,
load_dino_model,
load_hash_compressor,
load_sam_model,
)
from .proposal import (
extract_masked_region,
generate_proposals,
generate_proposals_batch,
)
def create_pipeline_from_config(config) -> "HashPipeline":
@@ -36,7 +36,6 @@ def create_pipeline_from_config(config) -> "HashPipeline":
sam_model=config.model.sam_model,
sam_min_mask_area=config.model.sam_min_mask_area,
sam_max_masks=config.model.sam_max_masks,
sam_points_per_batch=config.model.sam_points_per_batch,
hash_bits=config.model.compression_dim,
compressor_path=config.model.compressor_path,
)
@@ -60,7 +59,6 @@ class HashPipeline(nn.Module):
sam_model: str = "facebook/sam2.1-hiera-large",
sam_min_mask_area: int = 100,
sam_max_masks: int = 10,
sam_points_per_batch: int = 64,
hash_bits: int = 512,
compressor_path: Optional[str] = None,
):
@@ -69,15 +67,13 @@ class HashPipeline(nn.Module):
# Device for model placement.
self.device = get_device()
# SAM2 settings.
self.sam_model_name = sam_model
# SAM2 filter settings.
self.sam_min_mask_area = sam_min_mask_area
self.sam_max_masks = sam_max_masks
self.sam_points_per_batch = sam_points_per_batch
# Load models.
self.mask_generator = load_sam_model(model_name=sam_model)
self.processor, self.dino = load_dino_model(model_name=dino_model)
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.
self.dino_dim = get_dino_dim(dino_model)
@@ -94,25 +90,27 @@ class HashPipeline(nn.Module):
"""Number of bits in the hash code."""
return self.hash_compressor.hash_bits
def _segment_with_sam(self, image: Image.Image) -> Image.Image:
"""Segment image with SAM and extract the largest object mask.
If no valid masks are found, returns the original image.
def _segment_with_sam(
self, image: Image.Image, bboxes: list[list[float]]
) -> Image.Image:
"""Segment image with SAM and extract the best object mask.
Args:
image: Input PIL Image.
bboxes: Bounding boxes from object detector as [[x1,y1,x2,y2], ...].
Returns:
Masked image containing only the largest object, or original if no masks.
Masked image containing only the best object, or original if no masks.
"""
masks = generate_proposals(
self.mask_generator,
self.sam_model,
self.sam_processor,
image,
min_area=self.sam_min_mask_area,
max_masks=self.sam_max_masks,
points_per_batch=self.sam_points_per_batch,
bboxes,
)
masks = _filter_masks(masks, self.sam_min_mask_area, self.sam_max_masks)
if not masks:
return image
@@ -124,15 +122,19 @@ class HashPipeline(nn.Module):
def _segment_with_sam_dataset(
self,
images: Sequence[Image.Image],
bboxes_per_image: list[list[list[float]]],
) -> list[Image.Image]:
image_list = list(images)
masks_dataset = generate_proposals_batch(
self.mask_generator,
self.sam_model,
self.sam_processor,
image_list,
min_area=self.sam_min_mask_area,
max_masks=self.sam_max_masks,
points_per_batch=self.sam_points_per_batch,
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:
@@ -156,14 +158,14 @@ class HashPipeline(nn.Module):
Returns:
Last hidden state tokens of shape [1, N, dim].
"""
inputs = self.processor(image, return_tensors="pt").to(self.device)
inputs = self.dino_processor(image, return_tensors="pt").to(self.device)
with torch.no_grad():
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(
inputs = self.dino_processor(images=list(images), return_tensors="pt").to(
self.device
)
@@ -171,16 +173,17 @@ class HashPipeline(nn.Module):
outputs = self.dino(**inputs)
return outputs.last_hidden_state
def forward(self, image: Image.Image) -> torch.Tensor:
def forward(self, image: Image.Image, bboxes: list[list[float]]) -> torch.Tensor:
"""Process a single image through the full pipeline.
Args:
image: Input PIL Image.
bboxes: Bounding boxes from object detector as [[x1,y1,x2,y2], ...].
Returns:
Binary hash codes of shape [1, hash_bits] as int32.
"""
image = self._segment_with_sam(image)
image = self._segment_with_sam(image, bboxes)
tokens = self._dino_forward(image)
_, _, bits = self.hash_compressor(tokens)
return bits
@@ -188,6 +191,7 @@ class HashPipeline(nn.Module):
def forward_dataset(
self,
images: Sequence[Image.Image],
bboxes_per_image: list[list[list[float]]],
batch_size: int = 32,
apply_sam: bool = True,
) -> torch.Tensor:
@@ -202,7 +206,9 @@ class HashPipeline(nn.Module):
)
if apply_sam:
processed_images = self._segment_with_sam_dataset(image_list)
processed_images = self._segment_with_sam_dataset(
image_list, bboxes_per_image
)
else:
processed_images = image_list
@@ -215,16 +221,19 @@ class HashPipeline(nn.Module):
return torch.cat(batch_bits, dim=0)
def extract_features(self, image: Image.Image) -> torch.Tensor:
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)
image = self._segment_with_sam(image, bboxes)
tokens = self._dino_forward(image)
features = tokens.mean(dim=1)
return F.normalize(features, dim=-1)
@@ -232,6 +241,7 @@ class HashPipeline(nn.Module):
def extract_features_dataset(
self,
images: Sequence[Image.Image],
bboxes_per_image: list[list[list[float]]],
batch_size: int = 32,
apply_sam: bool = True,
) -> torch.Tensor:
@@ -246,7 +256,9 @@ class HashPipeline(nn.Module):
)
if apply_sam:
processed_images = self._segment_with_sam_dataset(image_list)
processed_images = self._segment_with_sam_dataset(
image_list, bboxes_per_image
)
else:
processed_images = image_list
@@ -258,3 +270,16 @@ class HashPipeline(nn.Module):
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]