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