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

# Load your trained YOLOv8 model
model = YOLO("YOLOV_8_model.pt")

def segment_conjunctiva(image):
    # Convert input image to RGB (if not already)
    if not isinstance(image, Image.Image):
        raise ValueError("Input is not a valid image")

    image = image.convert("RGB")  # Convert to RGB to avoid issues with mode
    # Save image to PNG format in memory (emulating PNG conversion)
    png_converted = Image.new("RGB", image.size)
    png_converted.paste(image)

    # Convert to numpy array for YOLO model
    image_np = np.array(png_converted)

    # Run prediction
    results = model(source=image_np, conf=0.5)

    # Check if masks are present (optional debug)
    if results[0].masks is None:
        print("⚠️ No segmentation masks detected")

    # Get annotated image with masks
    annotated_img = results[0].plot()

    return annotated_img

# Gradio Interface
title = "👁️Conjunctiva Segmentation using YOLOv8"
description = "Upload an eye image in any format (JPG/PNG/etc). This app uses a YOLOv8 segmentation model to segment the palpebral conjunctiva."

demo = gr.Interface(
    fn=segment_conjunctiva,
    inputs=gr.Image(type="pil"),
    outputs=gr.Image(type="numpy"),
    title=title,
    description=description,
    examples=[],
)

demo.launch()