Spaces:
Sleeping
Sleeping
File size: 1,036 Bytes
d1a1edf | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 | """SAM 2 Segmentation"""
from core.config import settings
import numpy as np
class SAM2Segmenter:
def __init__(self):
self.predictor = None
self.current_image = None
def load_model(self):
try:
from sam2.build_sam import build_sam2
from sam2.sam2_image_predictor import SAM2ImagePredictor
model = build_sam2("sam2_hiera_large.pt")
self.predictor = SAM2ImagePredictor(model)
except ImportError:
print("SAM 2 not installed. Segmentation disabled.")
def set_image(self, image):
self.current_image = image
if self.predictor:
self.predictor.set_image(image)
def segment_at_point(self, x: int, y: int) -> np.ndarray:
if not self.predictor:
return np.zeros((1, 1))
masks, scores, _ = self.predictor.predict(
point_coords=[[x, y]],
point_labels=[1],
multimask_output=False
)
return masks[0]
sam2_segmenter = SAM2Segmenter()
|