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

# Load the fine-tuned model
model = AutoModelForImageClassification.from_pretrained("Pavarissy/ConvNextV2-large-DogBreed")

# Initialize the image processor
preprocessor = AutoImageProcessor.from_pretrained("Pavarissy/ConvNextV2-large-DogBreed")

def classify_image(image):
    # Preprocess the image
    inputs = preprocessor(images=image, return_tensors="pt")

    # Model prediction
    with torch.no_grad():
        logits = model(**inputs).logits

    # Convert logits to probabilities
    probs = logits.softmax(dim=-1)

    # Extract top 5 predictions
    top_5_probs, top_5_labels = torch.topk(probs, 5)
    top_5_probs = top_5_probs.squeeze().tolist()
    top_5_labels = top_5_labels.squeeze().tolist()

    # Map labels to their names
    labels = model.config.id2label
    predicted_labels = [labels[label] for label in top_5_labels]

    return dict(zip(predicted_labels, top_5_probs))

# Create a Gradio interface
iface = gr.Interface(
    fn=classify_image,
    inputs=gr.Image(type="pil"),
    outputs=gr.Label(num_top_classes=5),
    title="Dog Breed Classifier",
    description="Upload an image of a dog, and the model will predict the breed."
)

# Launch the interface
iface.launch(share=True)