pevm/segmentation.py

import numpy as np
import torch
from ultralytics import YOLO

from pevm.setup import setup
from pevm.utils import resize_image, format_results, point_prompt

setup()
model = YOLO("./weights/FastSAM-x.pt")


def get_mask(raw_image, input_points=None, input_labels=None):
    if input_labels is None:
        input_labels = [1]
    if input_points is None:
        input_points = [[512, 512]]

    resized_image = resize_image(raw_image, input_size=1024)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    results = model(resized_image, device=device, retina_masks=True)
    results = format_results(results[0], 0)
    masks, _ = point_prompt(results, input_points, input_labels)

    return masks.astype(np.float32) * 255