| import numpy as np | |
| from sam2.build_sam import build_sam2 | |
| from sam2.sam2_image_predictor import SAM2ImagePredictor | |
| class Predictor: | |
| def __init__(self, model_cfg, checkpoint, device): | |
| self.device = device | |
| self.model = build_sam2(model_cfg, checkpoint, device=device) | |
| self.predictor = SAM2ImagePredictor(self.model) | |
| self.image_set = False | |
| def set_image(self, image): | |
| """Set the image for SAM prediction.""" | |
| self.image = image | |
| self.predictor.set_image(image) | |
| self.image_set = True | |
| def predict(self, point_coords, point_labels, multimask_output=False): | |
| """Run SAM prediction.""" | |
| if not self.image_set: | |
| raise RuntimeError("An image must be set with .set_image(...) before mask prediction.") | |
| return self.predictor.predict( | |
| point_coords=point_coords, | |
| point_labels=point_labels, | |
| multimask_output=multimask_output | |
| ) | |