From e815f70a3805ba9e7f0198b0ab1c193515d209ef Mon Sep 17 00:00:00 2001 From: Niels Date: Tue, 6 Aug 2024 08:32:36 +0200 Subject: [PATCH] Address comment --- sam2/build_sam.py | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/sam2/build_sam.py b/sam2/build_sam.py index 9bb5279..a061731 100644 --- a/sam2/build_sam.py +++ b/sam2/build_sam.py @@ -88,16 +88,22 @@ def build_sam2_hf(model_id, **kwargs): } 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) + 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 - config_file = hf_hub_download(repo_id=model_id, filename=f"{model_id}.yaml") - ckpt_path = hf_hub_download(repo_id=model_id, filename=f"{model_id}.pt") - return build_sam2_video_predictor(config_file=config_file, ckpt_path=ckpt_path, **kwargs) + 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):