from ultralytics import YOLO
import cv2
import torch

BEST_WEIGHT = './segmentation/weights/oai_s_best4.pt'
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")

def create_model(weight_path=None):
    if weight_path:
        return YOLO(weight_path)
    else:
        return YOLO(BEST_WEIGHT)

class Segmenter():
    def __init__(self, weight_path=None):
        self.model = create_model(weight_path).to(DEVICE)
    
    def segment(self, img):
        """
            input: image (H, W, C)
            output: mask (H, W) with femur is 1 and tibia is 2 
        """
        img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

        eimg = cv2.equalizeHist(img)

        clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
        eimg = clahe.apply(eimg)

        eimg = cv2.cvtColor(eimg, cv2.COLOR_GRAY2RGB)

        res = self.model(eimg, verbose=False)
        result = res[0]

        if result.masks is None or result.boxes is None:
            raise ValueError("Segmentation failed: no mask/object detected.")

        masks = result.masks.data
        classes = result.boxes.cls

        if masks is None or len(masks) < 2 or len(classes) < 2:
            raise ValueError("Segmentation failed: need at least 2 detected objects (femur and tibia).")

        mask = torch.zeros_like(masks[0])
        n = min(len(masks), len(classes))
        for i in range(n):
            class_value = int(classes[i].item()) + 1
            mask = torch.where(masks[i] > 0.5, torch.tensor(float(class_value), device=mask.device), mask)

        mask = mask.cpu().numpy()

        return mask