update to the latest sam2 version and support box prompts in video tracking
This commit is contained in:
@@ -76,6 +76,44 @@ def build_sam2_video_predictor(
|
||||
return model
|
||||
|
||||
|
||||
def build_sam2_hf(model_id, **kwargs):
|
||||
|
||||
from huggingface_hub import hf_hub_download
|
||||
|
||||
model_id_to_filenames = {
|
||||
"facebook/sam2-hiera-tiny": ("sam2_hiera_t.yaml", "sam2_hiera_tiny.pt"),
|
||||
"facebook/sam2-hiera-small": ("sam2_hiera_s.yaml", "sam2_hiera_small.pt"),
|
||||
"facebook/sam2-hiera-base-plus": (
|
||||
"sam2_hiera_b+.yaml",
|
||||
"sam2_hiera_base_plus.pt",
|
||||
),
|
||||
"facebook/sam2-hiera-large": ("sam2_hiera_l.yaml", "sam2_hiera_large.pt"),
|
||||
}
|
||||
config_name, checkpoint_name = model_id_to_filenames[model_id]
|
||||
ckpt_path = hf_hub_download(repo_id=model_id, filename=checkpoint_name)
|
||||
return build_sam2(config_file=config_name, ckpt_path=ckpt_path, **kwargs)
|
||||
|
||||
|
||||
def build_sam2_video_predictor_hf(model_id, **kwargs):
|
||||
|
||||
from huggingface_hub import hf_hub_download
|
||||
|
||||
model_id_to_filenames = {
|
||||
"facebook/sam2-hiera-tiny": ("sam2_hiera_t.yaml", "sam2_hiera_tiny.pt"),
|
||||
"facebook/sam2-hiera-small": ("sam2_hiera_s.yaml", "sam2_hiera_small.pt"),
|
||||
"facebook/sam2-hiera-base-plus": (
|
||||
"sam2_hiera_b+.yaml",
|
||||
"sam2_hiera_base_plus.pt",
|
||||
),
|
||||
"facebook/sam2-hiera-large": ("sam2_hiera_l.yaml", "sam2_hiera_large.pt"),
|
||||
}
|
||||
config_name, checkpoint_name = model_id_to_filenames[model_id]
|
||||
ckpt_path = hf_hub_download(repo_id=model_id, filename=checkpoint_name)
|
||||
return build_sam2_video_predictor(
|
||||
config_file=config_name, ckpt_path=ckpt_path, **kwargs
|
||||
)
|
||||
|
||||
|
||||
def _load_checkpoint(model, ckpt_path):
|
||||
if ckpt_path is not None:
|
||||
sd = torch.load(ckpt_path, map_location="cpu")["model"]
|
||||
|
Reference in New Issue
Block a user