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

First draft

Browse Source
This commit is contained in:
Niels
2024-08-03 12:57:05 +02:00
parent 57bc94b739
commit b72a8a97f0
2 changed files with 30 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
sam2/build_sam.py
View File

@@ -11,6 +11,8 @@ from hydra import compose
from hydra.utils import instantiate
from omegaconf import OmegaConf
from huggingface_hub import hf_hub_download
def build_sam2(
config_file,
@@ -76,6 +78,18 @@ def build_sam2_video_predictor(
return model
def build_sam2_hf(model_id, **kwargs):
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)
def build_sam2_video_predictor_hf(model_id, **kwargs):
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)
def _load_checkpoint(model, ckpt_path):
if ckpt_path is not None:
sd = torch.load(ckpt_path, map_location="cpu")["model"]
@@ -86,4 +100,4 @@ def _load_checkpoint(model, ckpt_path):
if unexpected_keys:
logging.error(unexpected_keys)
raise RuntimeError()
logging.info("Loaded checkpoint sucessfully")
logging.info("Loaded checkpoint sucessfully")