Spaces:
Sleeping
Sleeping
import torch | |
from transformers import SamModel, SamProcessor | |
class SAM(): | |
def __init__(self): | |
self.device = "cuda" if torch.cuda.is_available() else "cpu" | |
self.model = SamModel.from_pretrained("facebook/sam-vit-large").to(self.device) | |
self.processor = SamProcessor.from_pretrained("facebook/sam-vit-large") | |
def segment(self, raw_image, input_points): | |
inputs = self.processor(raw_image, input_points=input_points, return_tensors="pt").to(self.device) | |
with torch.no_grad(): | |
outputs = self.model(**inputs) | |
masks = self.processor.image_processor.post_process_masks( | |
outputs.pred_masks.cpu(), inputs["original_sizes"].cpu(), inputs["reshaped_input_sizes"].cpu() | |
) | |
scores = outputs.iou_scores | |
return masks, scores |