Quellcode durchsuchen

Address comment

Niels vor 1 Jahr
Ursprung
Commit
e815f70a38
1 geänderte Dateien mit 10 neuen und 4 gelöschten Zeilen
  1. 10 4
      sam2/build_sam.py

+ 10 - 4
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):