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

# Load your trained YOLOv8 model
model = YOLO('pytorch_yolov8_model.pt')  # Update this path to your model weights file

def detect_pneumonia(image):
    # YOLOv8 expects image in numpy array format (BGR)
    image_bgr = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)

    # Run prediction
    results = model.predict(image_bgr, conf=0.5)

    # Draw bounding boxes
    for r in results:
        for box in r.boxes:
            x1, y1, x2, y2 = map(int, box.xyxy[0])
            conf = box.conf[0]
            label = f'Pneumonia: {conf:.2f}'
            cv2.rectangle(image_bgr, (x1, y1), (x2, y2), (255, 0, 0), 2)
            cv2.putText(image_bgr, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (255, 0, 0), 2)

    # Convert back to RGB for display
    result_image = cv2.cvtColor(image_bgr, cv2.COLOR_BGR2RGB)
    return result_image

# Gradio Interface
interface = gr.Interface(
    fn=detect_pneumonia,
    inputs=gr.Image(type="pil"),
    outputs=gr.Image(type="pil"),
    title="Pneumonia Detection with YOLOv8",
    description="Upload a chest X-ray image. The model will predict pneumonia regions using bounding boxes."
)

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