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

feat(compressors): add object scoring and selection for SAM masks

Browse Source
This commit is contained in:
SikongJueluo
2026-03-30 16:31:06 +08:00
parent f6c1a67e88
commit a809803979
8 changed files with 581 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -7,6 +7,7 @@ import torch.nn as nn
import torch.nn.functional as F
from PIL import Image
from .object_score import select_best_mask
from utils import get_device
from utils.image import extract_masked_region, segment_image, segment_image_dataset
from utils.model import (
@@ -111,7 +112,10 @@ class HashPipeline(nn.Module):
if not masks:
return image
return extract_masked_region(image, masks[0]["segment"])
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,
@@ -125,10 +129,19 @@ class HashPipeline(nn.Module):
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)
]
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))
if best_mask is None:
selected_images.append(image)
continue
selected_images.append(extract_masked_region(image, best_mask["segment"]))
return selected_images
def _dino_forward(self, image: Image.Image) -> torch.Tensor:
"""Extract DINO tokens from an image.