# segmenter.py
from ultralytics import YOLO

class FoodSegmenter:
    def __init__(self, model_path="models/yolov8n-seg.pt"):
        self.model_path = model_path
        self.model = self._load_model()

    def _load_model(self):
        """
        Loads YOLOv8 segmentation model safely:
        - Falls back if 'device' argument is not supported.
        """
        try:
            # Newer YOLO versions may support 'device'
            model = YOLO(self.model_path, task="segment", device="cpu")
        except TypeError:
            # Older versions: skip 'device'
            model = YOLO(self.model_path, task="segment")
        return model

    def segment(self, image_path, conf=0.25):
        """
        Segment food items in an image.
        Returns YOLO results object.
        """
        results = self.model.predict(source=image_path, conf=conf, verbose=False)
        return results

# Example usage
if __name__ == "__main__":
    segmenter = FoodSegmenter()
    results = segmenter.segment("test_image.jpg")
    for r in results:
        print(r.masks)  # segmentation masks