import gradio as gr
from transformers import pipeline

# Load models
vit_classifier = pipeline("image-classification", model="Schram03/fruits-classification")
clip_detector = pipeline(model="openai/clip-vit-large-patch14", task="zero-shot-image-classification")

labels_fruits = [
    'Apple', 'Apricot', 'Avocado', 'Banana', 'Blackberry', 'Blueberry', 'Cactus fruit', 'Cherry', 'Clementine', 'Cocos', 'Grape', 'Grapefruit', 'Guava', 'Kaki', 'Kiwi', 'Lemon', 'Limes', 'Lychee', 'Mandarine', 'Mango', 'Maracuja', 'Nectarine', 'Orange', 'Papaya', 'Passion Fruit', 'Peach', 'Pear', 'Pineapple', 'Plum', 'Pomegranate', 'Quince', 'Strawberry', 'Watermelon'
]

def classify_fruit(image):
    vit_results = vit_classifier(image)
    vit_output = {result['label']: result['score'] for result in vit_results}
    
    clip_results = clip_detector(image, candidate_labels=labels_fruits)
    clip_output = {result['label']: result['score'] for result in clip_results}
    
    return {"ViT Classification": vit_output, "CLIP Zero-Shot Classification": clip_output}

example_images = [
    ["example_images/apple.jpg"],
    ["example_images/blueberry.jpg"],
    ["example_images/lemon.jpg"],
    ["example_images/pear.jpg"],
    ["example_images/plum.jpg"],
    ["example_images/watermelon.jpg"]
]

iface = gr.Interface(
    fn=classify_fruit,
    inputs=gr.Image(type="filepath"),
    outputs=gr.JSON(),
    title="Fruit Classification Comparison",
    description="Upload an image of a fruit, and compare results from a trained ViT model and a zero-shot CLIP model.",
    examples=example_images
)

iface.launch()