mirror of
https://github.com/SikongJueluo/Mini-Nav.git
synced 2026-07-12 20:15:31 +08:00
SikongJueluo
af0531a5eb
- switch OWLv2 loader return type to Owlv2ForObjectDetection - add detect_objects and detect_objects_batch with two-stage score filtering - add DetectionResult typed dict and conversion helper for post-processed outputs - export new detection APIs from proposal module
80 lines
2.1 KiB
Python
80 lines
2.1 KiB
Python
"""Model loading utilities for DINO, SAM2, OWLv2 and HashCompressor."""
|
|
|
|
from typing import TYPE_CHECKING, Any
|
|
|
|
import torch
|
|
from transformers import (
|
|
AutoImageProcessor,
|
|
AutoModel,
|
|
Owlv2ForObjectDetection,
|
|
Owlv2Processor,
|
|
Sam2Model,
|
|
Sam2Processor,
|
|
)
|
|
from utils import get_device
|
|
|
|
if TYPE_CHECKING:
|
|
from compressors.hash_compressor import HashCompressor
|
|
|
|
|
|
def load_sam_model(
|
|
model_name: str = "facebook/sam2.1-hiera-large",
|
|
) -> tuple[Sam2Processor, Sam2Model]:
|
|
"""Load SAM2 processor and model with frozen parameters."""
|
|
device = get_device()
|
|
|
|
processor = Sam2Processor.from_pretrained(model_name)
|
|
model = Sam2Model.from_pretrained(model_name).to(device)
|
|
model.eval()
|
|
for param in model.parameters():
|
|
param.requires_grad = False
|
|
|
|
return processor, model
|
|
|
|
|
|
def load_dino_model(
|
|
model_name: str = "facebook/dinov2-large",
|
|
) -> tuple[Any, Any]:
|
|
device = get_device()
|
|
|
|
processor = AutoImageProcessor.from_pretrained(model_name)
|
|
dino = AutoModel.from_pretrained(model_name).to(device)
|
|
dino.eval()
|
|
|
|
return processor, dino
|
|
|
|
|
|
def load_owlv2_model(
|
|
model_name: str = "google/owlv2-base-patch16-ensemble",
|
|
) -> tuple[Owlv2Processor, Owlv2ForObjectDetection]:
|
|
device = get_device()
|
|
|
|
processor = Owlv2Processor.from_pretrained(model_name)
|
|
model = Owlv2ForObjectDetection.from_pretrained(model_name).to(device)
|
|
model.eval()
|
|
|
|
return processor, model
|
|
|
|
|
|
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
|