import gradio as gr
import spaces
import torch
from transformers import AutoImageProcessor, AutoModelForImageClassification

MODEL_ID = "DiggaDee/brain-tumor-resnet18"

processor = AutoImageProcessor.from_pretrained(MODEL_ID)

model = AutoModelForImageClassification.from_pretrained(MODEL_ID)
model.to("cuda")
model.eval()


@spaces.GPU
def classify_mri(image):
    if image is None:
        return {}, "Please upload an MRI image."

    inputs = processor(
        images=image,
        return_tensors="pt"
    )

    inputs = {
        key: value.to("cuda")
        for key, value in inputs.items()
    }

    with torch.no_grad():
        outputs = model(**inputs)
        probabilities = torch.softmax(outputs.logits, dim=-1)[0]

    scores = {
        model.config.id2label[i]: float(probabilities[i])
        for i in range(len(probabilities))
    }

    prediction_index = probabilities.argmax().item()
    prediction = model.config.id2label[prediction_index]
    confidence = probabilities[prediction_index].item()

    result = (
        f"### Prediction: {prediction}\n\n"
        f"Confidence: **{confidence:.1%}**\n\n"
        "Educational prototype only. Not for medical diagnosis."
    )

    return scores, result


with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown(
        """
        # Brain MRI Tumour Classification

        Upload a brain MRI image to classify it as:

        **Glioma · Meningioma · No Tumour · Pituitary**
        """
    )

    with gr.Row():
        with gr.Column():
            image_input = gr.Image(
                type="pil",
                label="Upload MRI Image"
            )

            classify_button = gr.Button(
                "Classify MRI",
                variant="primary"
            )

        with gr.Column():
            prediction_output = gr.Label(
                label="Class Probabilities",
                num_top_classes=4
            )

            result_output = gr.Markdown()

    classify_button.click(
        fn=classify_mri,
        inputs=image_input,
        outputs=[
            prediction_output,
            result_output
        ]
    )

    gr.Markdown(
        """
        **Disclaimer:** This application is for educational purposes only
        and must not be used for clinical diagnosis.
        """
    )


demo.launch()