| import numpy as np | |
| import supervision as sv | |
| from segment_anything.modeling.sam import Sam | |
| from segment_anything import SamPredictor, SamAutomaticMaskGenerator | |
| def sam_inference( | |
| image: np.ndarray, | |
| model: Sam | |
| ) -> sv.Detections: | |
| mask_generator = SamAutomaticMaskGenerator(model) | |
| result = mask_generator.generate(image=image) | |
| return sv.Detections.from_sam(result) | |
| def sam_interactive_inference( | |
| image: np.ndarray, | |
| mask: np.ndarray, | |
| model: Sam | |
| ) -> sv.Detections: | |
| predictor = SamPredictor(model) | |
| predictor.set_image(image) | |
| masks = [] | |
| for polygon in sv.mask_to_polygons(mask.astype(bool)): | |
| random_point_indexes = np.random.choice(polygon.shape[0], size=5, replace=True) | |
| input_point = polygon[random_point_indexes] | |
| input_label = np.ones(5) | |
| mask = predictor.predict( | |
| point_coords=input_point, | |
| point_labels=input_label, | |
| multimask_output=False, | |
| )[0][0] | |
| masks.append(mask) | |
| masks = np.array(masks, dtype=bool) | |
| return sv.Detections( | |
| xyxy=sv.mask_to_xyxy(masks), | |
| mask=masks | |
| ) | |