mirror of
https://github.com/SikongJueluo/Mini-Nav.git
synced 2026-07-12 20:15:31 +08:00
58 lines
1.5 KiB
Python
58 lines
1.5 KiB
Python
"""Model loading utilities for DINO, SAM2 and HashCompressor."""
|
|
|
|
from compressors import HashCompressor
|
|
|
|
|
|
import torch
|
|
from transformers import AutoImageProcessor, AutoModel, pipeline, MaskGenerationPipeline
|
|
|
|
from .common import get_device
|
|
|
|
|
|
def load_sam_model(
|
|
model_name: str = "facebook/sam2.1-hiera-large",
|
|
) -> MaskGenerationPipeline:
|
|
device = str(get_device())
|
|
device_id = 0 if device.startswith("cuda") else -1
|
|
|
|
return pipeline(
|
|
task="mask-generation",
|
|
model=model_name,
|
|
device=device_id,
|
|
)
|
|
|
|
|
|
def load_dino_model(
|
|
model_name: str = "facebook/dinov2-large",
|
|
) -> tuple[AutoImageProcessor, AutoModel]:
|
|
device = get_device()
|
|
|
|
processor = AutoImageProcessor.from_pretrained(model_name)
|
|
dino = AutoModel.from_pretrained(model_name).to(device)
|
|
dino.eval()
|
|
|
|
return processor, dino
|
|
|
|
|
|
def get_dino_dim(model_name: str) -> int:
|
|
if "large" in model_name.lower():
|
|
return 1024
|
|
return 768
|
|
|
|
|
|
def load_hash_compressor(
|
|
input_dim: int = 1024,
|
|
hash_bits: int = 512,
|
|
compressor_path: str | None = None,
|
|
) -> HashCompressor:
|
|
from compressors.hash_compressor import HashCompressor
|
|
|
|
device = get_device()
|
|
compressor = HashCompressor(input_dim=input_dim, hash_bits=hash_bits).to(device)
|
|
|
|
if compressor_path is not None:
|
|
compressor.load_state_dict(torch.load(compressor_path, map_location=device))
|
|
print(f"[OK] Loaded HashCompressor from {compressor_path}")
|
|
|
|
return compressor
|