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

Format using ufmt

Browse Source
This commit is contained in:
Niels
2024-08-06 22:43:35 +02:00
parent 27a167c004
commit 8f15c6255a
2 changed files with 13 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
sam2/build_sam.py
View File

@@ -83,7 +83,10 @@ def build_sam2_hf(model_id, **kwargs):
model_id_to_filenames = { model_id_to_filenames = {
"facebook/sam2-hiera-tiny": ("sam2_hiera_t.yaml", "sam2_hiera_tiny.pt"), "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-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-base-plus": (
"sam2_hiera_b+.yaml",
"sam2_hiera_base_plus.pt",
),
"facebook/sam2-hiera-large": ("sam2_hiera_l.yaml", "sam2_hiera_large.pt"), "facebook/sam2-hiera-large": ("sam2_hiera_l.yaml", "sam2_hiera_large.pt"),
} }
config_name, checkpoint_name = model_id_to_filenames[model_id] config_name, checkpoint_name = model_id_to_filenames[model_id]
@@ -98,12 +101,17 @@ def build_sam2_video_predictor_hf(model_id, **kwargs):
model_id_to_filenames = { model_id_to_filenames = {
"facebook/sam2-hiera-tiny": ("sam2_hiera_t.yaml", "sam2_hiera_tiny.pt"), "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-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-base-plus": (
"sam2_hiera_b+.yaml",
"sam2_hiera_base_plus.pt",
),
"facebook/sam2-hiera-large": ("sam2_hiera_l.yaml", "sam2_hiera_large.pt"), "facebook/sam2-hiera-large": ("sam2_hiera_l.yaml", "sam2_hiera_large.pt"),
} }
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_video_predictor(
config_file=config_name, ckpt_path=ckpt_path, **kwargs
)
def _load_checkpoint(model, ckpt_path): def _load_checkpoint(model, ckpt_path):
@@ -116,4 +124,4 @@ def _load_checkpoint(model, ckpt_path):
if unexpected_keys: if unexpected_keys:
logging.error(unexpected_keys) logging.error(unexpected_keys)
raise RuntimeError() raise RuntimeError()
logging.info("Loaded checkpoint sucessfully") logging.info("Loaded checkpoint sucessfully")

2
sam2/sam2_video_predictor.py
View File

@@ -119,7 +119,7 @@ class SAM2VideoPredictor(SAM2Base):
sam_model = build_sam2_video_predictor_hf(model_id, **kwargs) sam_model = build_sam2_video_predictor_hf(model_id, **kwargs)
return cls(sam_model) return cls(sam_model)
def _obj_id_to_idx(self, inference_state, obj_id): def _obj_id_to_idx(self, inference_state, obj_id):
"""Map client-side object id to model-side object index.""" """Map client-side object id to model-side object index."""
obj_idx = inference_state["obj_id_to_idx"].get(obj_id, None) obj_idx = inference_state["obj_id_to_idx"].get(obj_id, None)