import gradio as gr
from transformers import AutoModelForImageClassification, AutoTokenizer
from PIL import Image

# Load Hugging Face model and tokenizer
model_name = 'ImageDifferentiator.pkl'  # Replace with the specific model name
model = AutoModelForImageClassification.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

# Get class labels from the model configuration
labels = model.config.id2label

def predict(img):
    # Tokenize and preprocess the image
    inputs = tokenizer(img, return_tensors="pt")

    # Make prediction using the Hugging Face model
    outputs = model(**inputs)
    logits = outputs.logits

    # Get class probabilities
    probs = torch.nn.functional.softmax(logits, dim=-1)[0].tolist()

    # Create result dictionary
    return {labels[i]: float(probs[i]) for i in range(len(labels))}

iface = gr.Interface(
    fn=predict,
    inputs=gr.inputs.Image(shape=(512, 512)),
    outputs=gr.outputs.Label(num_top_classes=3)
)

iface.launch(share=True)