quoccuongle/Grounded-SAM-2
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Address comment

Browse Source
This commit is contained in:
Niels
2024-08-06 08:32:36 +02:00
parent fbf7e3a664
commit e815f70a38
1 changed files with 10 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
sam2/build_sam.py
View File

@@ -88,16 +88,22 @@ def build_sam2_hf(model_id, **kwargs):
} }
config_name, checkpoint_name = model_id_to_filenames[model_id] config_name, checkpoint_name = model_id_to_filenames[model_id]
ckpt_path = hf_hub_download(repo_id=model_id, filename=checkpoint_name) 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): def build_sam2_video_predictor_hf(model_id, **kwargs):
from huggingface_hub import hf_hub_download from huggingface_hub import hf_hub_download
config_file = hf_hub_download(repo_id=model_id, filename=f"{model_id}.yaml") model_id_to_filenames = {
ckpt_path = hf_hub_download(repo_id=model_id, filename=f"{model_id}.pt") "facebook/sam2-hiera-tiny": ("sam2_hiera_t.yaml", "sam2_hiera_tiny.pt"),
return build_sam2_video_predictor(config_file=config_file, ckpt_path=ckpt_path, **kwargs) "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): def _load_checkpoint(model, ckpt_path):