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

# Load the YOLOv8 model
model = YOLO("yolov8n.pt")  # Replace with your trained brain tumor model

def predict(image_path):
    # Run YOLOv8 inference
    results = model(image_path)
    
    # Get annotated image
    annotated_frame = results[0].plot()
    
    # Convert BGR to RGB
    annotated_frame_rgb = cv2.cvtColor(annotated_frame, cv2.COLOR_BGR2RGB)
    
    # Check if a tumor is detected
    tumor_detected = len(results[0].boxes) > 0
    detection_message = "Tumor Detected!" if tumor_detected else "No Tumor Detected."
    
    return annotated_frame_rgb, detection_message

# Create Gradio interface
interface = gr.Interface(
    fn=predict,
    inputs=gr.Image(type="filepath", label="Upload MRI Image"),
    outputs=[
        gr.Image(label="Detection Result"),
        gr.Textbox(label="Diagnosis")
    ],
    title="Brain Tumor Detection with YOLOv8",
    description="Upload an MRI scan to detect brain tumors using AI.",
    allow_flagging="never"
)

interface.launch()