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

# Load YOLO model
model = YOLO("best.pt")  # Ensure 'best.pt' is in the same directory

# Define prediction function
def predict(image):
    # Perform YOLO detection
    result = model.predict(source=image, imgsz=640, conf=0.25)
    annotated_image = result[0].plot()

    # Convert image from BGR to RGB
    annotated_image = cv2.cvtColor(annotated_image, cv2.COLOR_BGR2RGB)

    return annotated_image

# Gradio interface
app = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="numpy", label="Upload an Image"),
    outputs=gr.Image(type="numpy", label="Detected Tooth Cavity"),
    title="Tooth Cavity Detection Using YOLO V10 by Pulastya 😎",
    description="Upload a dental Photo, and the YOLO V10 model will detect and annotate tooth decay."
)

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