Schram03/fruits-classification

1import gradio as gr2from transformers import pipeline3 4# Load models5vit_classifier = pipeline("image-classification", model="Schram03/fruits-classification")6clip_detector = pipeline(model="openai/clip-vit-large-patch14", task="zero-shot-image-classification")7 8labels_fruits = [9    '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'10]11 12def classify_fruit(image):13    vit_results = vit_classifier(image)14    vit_output = {result['label']: result['score'] for result in vit_results}15    16    clip_results = clip_detector(image, candidate_labels=labels_fruits)17    clip_output = {result['label']: result['score'] for result in clip_results}18    19    return {"ViT Classification": vit_output, "CLIP Zero-Shot Classification": clip_output}20 21example_images = [22    ["example_images/apple.jpg"],23    ["example_images/blueberry.jpg"],24    ["example_images/lemon.jpg"],25    ["example_images/pear.jpg"],26    ["example_images/plum.jpg"],27    ["example_images/watermelon.jpg"]28]29 30iface = gr.Interface(31    fn=classify_fruit,32    inputs=gr.Image(type="filepath"),33    outputs=gr.JSON(),34    title="Fruit Classification Comparison",35    description="Upload an image of a fruit, and compare results from a trained ViT model and a zero-shot CLIP model.",36    examples=example_images37)38 39iface.launch()