from ultralytics import YOLO
from PIL import Image
import torch
import io
import base64
import numpy as np
import cv2
from huggingface_hub import InferenceMixin

class Model:
    def __init__(self):
        # Load all your gingivitis models here
        self.model_redness = YOLO("best_redness2.pt")
        self.model_swelling = YOLO("best_swelling.pt")
        self.model_bleeding = YOLO("bleeding2.pt")

    def predict(self, image: Image.Image):
        # Convert image to OpenCV format
        img = np.array(image)

        # Run detection (choose one for now)
        results = self.model_redness(img)

        # Annotate the result image
        annotated = results[0].plot()  # returns a NumPy array (with boxes)
        _, buffer = cv2.imencode(".png", annotated)
        img_bytes = buffer.tobytes()

        # Return base64 image string
        img_base64 = base64.b64encode(img_bytes).decode("utf-8")
        return {"image_base64": img_base64}

model = Model()

def predict(image: Image.Image):
    return model.predict(image)