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
    # Get the predicted class
    predicted_class = logits.argmax(-1).item()
    # You may need to adjust the following line based on your class labels
    class_names = ["class1", "class2", "class3", "class4"]
    return predicted_class

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