import gradio as gr
from PIL import Image
from pathlib import Path
import numpy as np
from ultralytics import YOLO

MODEL_WEIGHTS_PATH = Path("weights/best.pt")
VERSION_PATH = Path("VERSION")

# Read version string from VERSION file
try:
    VERSION = VERSION_PATH.read_text().strip()
except Exception:
    VERSION = "unknown"

model = None
def get_model() -> YOLO:
    """
    Returns the YOLO model instance.
    """
    global model
    if model is None:
        if not MODEL_WEIGHTS_PATH.exists():
            raise FileNotFoundError(f"Model weights not found at {MODEL_WEIGHTS_PATH}. Please deploy weights before running.")
        model = YOLO(str(MODEL_WEIGHTS_PATH))
    return model

def segment(image: Image.Image) -> tuple[Image.Image, str]:
    """
    Returns a tuple: (segmentation mask PIL.Image, model version string)
    """
    model = get_model()
    img_np = np.array(image)
    results = model(img_np)
    if not results or not hasattr(results[0], "masks") or results[0].masks is None:
        mask_img = Image.new("L", image.size, 0)
    else:
        mask = results[0].masks.data[0].cpu().numpy()
        mask_img = Image.fromarray((mask * 255).astype(np.uint8))
        mask_img = mask_img.resize(image.size)
    return mask_img, str(VERSION)

iface = gr.Interface(
    fn=segment,
    inputs=gr.Image(type="pil"),
    outputs=[gr.Image(type="pil", label="Segmentation Mask"), gr.Textbox(label="Model Version")],
    title=f"YOLO Segmentation Model (version: {VERSION})",
    description=f"Upload an image to get a segmentation mask. Model version: {VERSION}"
)

if __name__ == "__main__":
    iface.launch()