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

# Load your YOLO model (replace with your model path)
model = YOLO("MPOX.pt")

def detect_mpox(image):
    # Run YOLO prediction
    results = model.predict(source=image, conf=0.25)

    # Extract detections
    detections = results[0].boxes
    num_detected = len(detections)

    # Draw bounding boxes on image
    img = np.array(image)
    for box in detections:
        x1, y1, x2, y2 = map(int, box.xyxy[0])
        confidence = float(box.conf[0])
        cls = int(box.cls[0])
        
        # Draw rectangle and label
        cv2.rectangle(img, (x1, y1), (x2, y2), (0, 0, 255), 2)
        label = f"MPOX {confidence:.2f}"
        cv2.putText(img, label, (x1, y1 - 10),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 0, 255), 2)

    # Confirmed case if any detection
    status = "Confirmed Case" if num_detected > 2 else "Not Confirmed"

    return img, f"Number of MPOX detected: {num_detected}\nStatus: {status}"

# Build Gradio interface
interface = gr.Interface(
    fn=detect_mpox,
    inputs=gr.Image(type="pil", label="Upload Body Image"),
    outputs=[
        gr.Image(type="numpy", label="Detection Result"),
        gr.Textbox(label="Detection Summary")
    ],
    title="MPOX Detection Prototype",
    description="Upload a body image to detect and visualize MPOX lesions using YOLO."
)

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