import gradio as gr
from transformers import AutoFeatureExtractor, AutoModelForImageClassification

# Load the feature extractor and model directly
feature_extractor = AutoFeatureExtractor.from_pretrained("Devarshi/Brain_Tumor_Classification")
model = AutoModelForImageClassification.from_pretrained("Devarshi/Brain_Tumor_Classification")

# Define the prediction function using the loaded model
def classify_image(image):
    # Preprocess the image and obtain features
    inputs = feature_extractor(images=image, return_tensors="pt")
    # Make prediction using the model
    outputs = model(**inputs)
    logits = outputs.logits
    # Get the predicted class and confidence of the prediction
    predicted_class = logits.argmax(dim=1).item()
    confidence = logits.softmax(dim=1).max().item()

    # Map the predicted class to the correct names
    class_names = ["glioma_tumor", "meningioma_tumor", "no_tumor", "pituitary_tumor"]
    predicted_class_text = class_names[predicted_class]

    return {"prediction": predicted_class_text, "confidence": confidence}

# Define the Gradio interface
iface = gr.Interface(
    fn=classify_image,
    inputs=gr.Image(),
    outputs="json",
    title="Brain Tumor Image Classification",
    description="This app classifies images of brain tumors into different classes.",
)

# Launch the Gradio interface
if __name__ == "__main__":
    iface.launch()