File size: 1,654 Bytes
43abac3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 | import os
import urllib.request
import torch
from segment_anything import sam_model_registry, SamAutomaticMaskGenerator
def _download_sam_checkpoint(checkpoint_path="outputs/models/sam_vit_b_01ec64.pth"):
"""
Downloads the Segment Anything Model (SAM) ViT-B checkpoint (350MB).
Args:
checkpoint_path (str): Local path to save/check the .pth file.
"""
url = "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth"
if not os.path.exists(checkpoint_path):
print(f"Downloading SAM checkpoint to {checkpoint_path}...")
os.makedirs(os.path.dirname(checkpoint_path), exist_ok=True)
urllib.request.urlretrieve(url, checkpoint_path)
print("Download complete.")
else:
print("SAM checkpoint already exists.")
def get_sam_generator(model_type="vit_b", checkpoint_path="outputs/models/sam_vit_b_01ec64.pth", device="cuda"):
"""
Initializes the SAM model and returns an automatic mask generator.
Args:
model_type (str): The SAM architecture version (e.g., 'vit_h').
checkpoint_path (str): Path to the model weights.
device (str): Device to run the model on ('cuda' or 'cpu').
Returns:
SamAutomaticMaskGenerator: Initialized mask generator object.
"""
if device == "cuda" and not torch.cuda.is_available():
device = "cpu"
print("CUDA not available. Falling back to CPU for SAM.")
_download_sam_checkpoint(checkpoint_path)
sam = sam_model_registry[model_type](checkpoint=checkpoint_path)
sam.to(device=device)
mask_generator = SamAutomaticMaskGenerator(sam)
return mask_generator
|