import gradio as gr
from utils import load_model, predict_image

# Load model once
model = load_model("model/best_classification_model.pth", num_classes=3)
class_names = ['COVID', 'Normal', 'Viral Pneumonia']

def classify_xray(img):
    prediction = predict_image(img, model, class_names)
    return f"🧠 Predicted: {prediction}"

# Gradio UI
title = "🩻 COVID-19 Chest X-ray Classifier"
description = """
Upload a Chest X-ray image and let the AI classify it as:
- **COVID**
- **Normal**
- **Viral Pneumonia**
<br><br>
🧠 Powered by ResNet18 + PyTorch
"""

demo = gr.Interface(
    fn=classify_xray,
    inputs=gr.Image(type="filepath", label="Upload Chest X-ray"),
    outputs=gr.Textbox(label="Prediction"),
    title=title,
    description=description,
    examples=["examples/x-ray.jpg"],
    theme="soft",  # use "huggingface" or "default" if preferred
)

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