mirror of
				https://github.com/facebookresearch/sam2.git
				synced 2025-11-04 11:32:12 +08:00 
			
		
		
		
	Make huggingface_hub soft dependency
This commit is contained in:
		
							parent
							
								
									0c28c630c2
								
							
						
					
					
						commit
						6aeee34775
					
				@ -11,8 +11,6 @@ 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,
 | 
			
		||||
@ -80,6 +78,8 @@ def build_sam2_video_predictor(
 | 
			
		||||
 | 
			
		||||
def build_sam2_hf(model_id, **kwargs):
 | 
			
		||||
 | 
			
		||||
    from huggingface_hub import hf_hub_download
 | 
			
		||||
 | 
			
		||||
    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"),
 | 
			
		||||
@ -87,12 +87,14 @@ def build_sam2_hf(model_id, **kwargs):
 | 
			
		||||
        "facebook/sam2-hiera-large": ("sam2_hiera_l.yaml", "sam2_hiera_large.pt"),
 | 
			
		||||
    }
 | 
			
		||||
    config_name, checkpoint_name = model_id_to_filenames[model_id]
 | 
			
		||||
    # config_file = hf_hub_download(repo_id=model_id, filename=config_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)
 | 
			
		||||
 | 
			
		||||
 | 
			
		||||
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)
 | 
			
		||||
 | 
			
		||||
@ -13,7 +13,6 @@ import torch
 | 
			
		||||
from PIL.Image import Image
 | 
			
		||||
 | 
			
		||||
from sam2.modeling.sam2_base import SAM2Base
 | 
			
		||||
from sam2.build_sam import build_sam2_hf
 | 
			
		||||
from sam2.utils.transforms import SAM2Transforms
 | 
			
		||||
 | 
			
		||||
 | 
			
		||||
@ -74,6 +73,8 @@ class SAM2ImagePredictor:
 | 
			
		||||
        Returns:
 | 
			
		||||
          (SAM2ImagePredictor): The loaded model.
 | 
			
		||||
        """
 | 
			
		||||
        from sam2.build_sam import build_sam2_hf
 | 
			
		||||
 | 
			
		||||
        sam_model = build_sam2_hf(model_id, **kwargs)
 | 
			
		||||
        return cls(sam_model)
 | 
			
		||||
 | 
			
		||||
 | 
			
		||||
@ -103,6 +103,23 @@ class SAM2VideoPredictor(SAM2Base):
 | 
			
		||||
        self._get_image_feature(inference_state, frame_idx=0, batch_size=1)
 | 
			
		||||
        return inference_state
 | 
			
		||||
 | 
			
		||||
    @classmethod
 | 
			
		||||
    def from_pretrained(cls, model_id: str, **kwargs) -> "SAM2ImagePredictor":
 | 
			
		||||
        """
 | 
			
		||||
        Load a pretrained model from the Hugging Face model hub.
 | 
			
		||||
 | 
			
		||||
        Arguments:
 | 
			
		||||
          model_id (str): The Hugging Face repository ID.
 | 
			
		||||
          **kwargs: Additional arguments to pass to the model constructor.
 | 
			
		||||
 | 
			
		||||
        Returns:
 | 
			
		||||
          (SAM2ImagePredictor): The loaded model.
 | 
			
		||||
        """
 | 
			
		||||
        from sam2.build_sam import build_sam2_video_predictor_hf
 | 
			
		||||
 | 
			
		||||
        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)
 | 
			
		||||
 | 
			
		||||
		Loading…
	
	
			
			x
			
			
		
	
		Reference in New Issue
	
	Block a user