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

# Load the model with original name
model = YOLO("best (2).pt")

def detect_teeth(image):
    results = model.predict(image)

    r = results[0]
    if r.masks is not None:
        r.masks.data = (r.masks.data > 0.3).float()

    # Render without confidence scores or labels
    output_img = r.plot(conf=False, labels=True, boxes=True)
    return Image.fromarray(output_img)

# Gradio interface
interface = gr.Interface(
    fn=detect_teeth,
    inputs=gr.Image(type="pil", label="Upload Dental X-ray"),
    outputs=gr.Image(type="pil", label="Detection Result"),
    title="Dental Segmentation App",
    description="Upload a dental X-ray to see the segmentation result using YOLOv8."
)

interface.launch()