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

MODEL_WEIGHTS_PATH = Path("weights/best.pt")  # Path to model weights (populated by deploy.sh)
VERSION_PATH = Path("VERSION")

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

# Lazy-load model (singleton)
model = None
def get_model():
    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):
    model = get_model()
    img_np = np.array(image)
    # Run prediction
    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)  # Blank mask if no detections
    else:
        mask = results[0].masks.data[0].cpu().numpy()  # (H, W) binary mask
        mask_img = Image.fromarray((mask * 255).astype(np.uint8))
        mask_img = mask_img.resize(image.size)  # Ensure mask matches input size
    # Return both the mask and version in the API response
    return {"mask": mask_img, "version": 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()