Niels преди 1 година
родител
ревизия
8f15c6255a
променени са 2 файла, в които са добавени 13 реда и са изтрити 5 реда
  1. 12 4
      sam2/build_sam.py
  2. 1 1
      sam2/sam2_video_predictor.py

+ 12 - 4
sam2/build_sam.py

@@ -83,7 +83,10 @@ def build_sam2_hf(model_id, **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-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]
@@ -98,12 +101,17 @@ def build_sam2_video_predictor_hf(model_id, **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-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)
+    return build_sam2_video_predictor(
+        config_file=config_name, ckpt_path=ckpt_path, **kwargs
+    )
 
 
 def _load_checkpoint(model, ckpt_path):
@@ -116,4 +124,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")

+ 1 - 1
sam2/sam2_video_predictor.py

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