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

# Model load karein (ensure karein best.pt upload ho chuki hai)
model = YOLO("best.pt")

def predict_image(img):
    # Model prediction
    results = model(img)
    
    # Annotated image hasil karein (boxes ke saath)
    # results[0].plot() humein numpy array deta hai jis par boxes bane hote hain
    res_plotted = results[0].plot()
    
    # RGB mein convert karein taake Gradio sahi dikhaye
    output_img = Image.fromarray(res_plotted[:, :, ::-1]) 
    
    return output_img

# Gradio interface setup
demo = gr.Interface(
    fn=predict_image,
    inputs=gr.Image(type="pil"),
    outputs=gr.Image(type="pil"),
    title="Face Mask Detection (YOLOv8)",
    description="Apni photo upload karein. Model 'with_mask', 'without_mask', ya 'incorrect' detect karega."
)

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