File size: 583 Bytes
811ffe8 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 |
from transformers import SamModel, SamProcessor
from PIL import Image
segmentation_model_id = "facebook/sam-vit-base"
sam_processor = SamProcessor.from_pretrained(segmentation_model_id)
sam_model = SamModel.from_pretrained(segmentation_model_id)
def run_segmentation(image_path: str):
"""SAM: return binary masks as nested lists"""
img = Image.open(image_path).convert("RGB")
inputs = sam_processor(images=img, return_tensors="pt")
outputs = sam_model(**inputs)
masks = outputs.pred_masks.squeeze(0).cpu().detach().numpy().tolist()
return {"masks": masks} |