423A35C7/sam2
1
0
Fork 0
mirror of https://github.com/facebookresearch/sam2.git synced 2025-09-18 04:32:48 +08:00
Code Issues Packages Projects Releases Wiki Activity

Use classmethod

Browse Source
This commit is contained in:
Niels 2024-08-03 14:14:12 +02:00
parent b72a8a97f0
commit 17b74501fb
1 changed files with 3 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
sam2/sam2_image_predictor.py
View File

@ -62,7 +62,8 @@ class SAM2ImagePredictor:
(64, 64),
]
def from_pretrained(model_id: str, **kwargs) -> "SAM2ImagePredictor":
@classmethod
def from_pretrained(cls, model_id: str, **kwargs) -> "SAM2ImagePredictor":
"""
Load a pretrained model from the Hugging Face model hub.
@ -74,7 +75,7 @@ class SAM2ImagePredictor:
(SAM2ImagePredictor): The loaded model.
"""
sam_model = build_sam2_hf(model_id, **kwargs)
return SAM2ImagePredictor(sam_model)
return cls(sam_model)
@torch.no_grad()
def set_image(