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

model = YOLO('best_V4.pt')

def predict(image):
    
    results = model(image, conf=0.8)
    
    detected = False
    LABEL_MAP = {
    0: "Other",
    1: "Pneumonia"
    }
    
    labels_found = []
    for result in results:
        boxes = result.boxes.xyxy.cpu().numpy()
        confidences = result.boxes.conf.cpu().numpy()
        class_ids = result.boxes.cls.cpu().numpy()
        
        for box, confidence, class_id in zip(boxes, confidences, class_ids):
            x1, y1, x2, y2 = map(int, box[:4])
            label = LABEL_MAP.get(int(class_id), "Unknown")
            cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2)
            label_text = f"{label} {confidence:.2f}"
            labels_found.append(label_text)
            
            cv2.putText(image, label_text, (x1, y1 - 10), 
                        cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
            detected = True 

    if not detected:
        return None, "No detected"

    pil_image = Image.fromarray(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
    message = "\n".join(labels_found)
    return pil_image, message

demo = gr.Interface(
    fn=predict, 
    inputs=gr.Image(type="numpy"), 
    outputs=[
        gr.Image(type="pil", label="Detection Result"), 
        gr.Textbox(label="Message")
    ],
    allow_flagging="never",
    api_name=False  # <- THIS DISABLES OPENAPI GENERATION AND AVOIDS THE ERROR
)

demo.launch(share=True, debug=True)