import gradio as gr
from transformers import AutoModelForImageClassification, AutoFeatureExtractor
import torch
from PIL import Image

# Load the model and feature extractor once during initialization
model_name = "amjadfqs/finalProject"
model = AutoModelForImageClassification.from_pretrained(model_name)
feature_extractor = AutoFeatureExtractor.from_pretrained(model_name)

def predict(image):
    # Preprocess the image
    inputs = feature_extractor(images=image, return_tensors="pt")
    # Make prediction
    with torch.no_grad():
        outputs = model(**inputs)
    logits = outputs.logits
    # Calculate the confidence values
    softmax = torch.nn.functional.softmax(logits, dim=1)
    confidences = softmax.squeeze().tolist()
    # Get the predicted class
    predicted_class_index = logits.argmax(-1).item()
    class_names = ["glioma", "meningioma", "notumor", "pituitary"]
    predicted_class = class_names[predicted_class_index]
    # Create a dictionary to return both the predicted class and the confidence values
    result = {
        "predicted_class": predicted_class,
        "confidences": {class_names[i]: confidences[i] for i in range(len(class_names))}
    }
    return result

# Set up the Gradio interface
image_cp = gr.Image(type="pil", label='Brain')
interface = gr.Interface(fn=predict, inputs=image_cp, outputs="json")
interface.launch()