| 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 |
|
|
