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

# Load a pre-trained image classification model
model_name = "Shio-Koube/Anime_filterer"
image_processor = AutoImageProcessor.from_pretrained(model_name)
model = AutoModelForImageClassification.from_pretrained(model_name)

def classify_image(image):
    # Ensure the image is in RGB mode
    if image is None:
        return "No image uploaded"
    
    # Convert image to RGB if needed
    if image.mode != 'RGB':
        image = image.convert('RGB')
    
    # Preprocess the image
    inputs = image_processor(images=image, return_tensors="pt")
    
    # Perform prediction
    with torch.no_grad():
        outputs = model(**inputs)
        logits = outputs.logits
    
    # Get predictions
    probabilities = torch.nn.functional.softmax(logits, dim=-1)
    
    # Get class labels and handle fewer than 5 classes
    labels = model.config.id2label
    num_classes = len(labels)
    
    # Determine number of predictions to show
    top_k = min(num_classes, 3)
    
    # Get top predictions
    top_prob, top_indices = probabilities.topk(top_k)
    
    # Format results
    results = []
    for prob, idx in zip(top_prob[0], top_indices[0]):
        label = labels[idx.item()]
        percentage = prob.item() * 100
        results.append(f"{label}: {percentage:.2f}%")
    
    return "\n".join(results)

# Create Gradio interface
iface = gr.Interface(
    fn=classify_image,
    inputs=gr.Image(type="pil"),
    outputs=gr.Textbox(label="Top Predictions"),
    title="Image Classification with Hugging Face",
    description="Upload an image to get classification predictions"
)

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