import gradio as gr
from transformers import pipeline

vit_classifier = pipeline("image-classification", model="maceythm/vit-90-animals")
clip_detector = pipeline(model="openai/clip-vit-large-patch14", task="zero-shot-image-classification")

labels_animals = [
    'antelope', 'badger', 'bat', 'bear', 'bee', 'beetle', 'bison', 'boar', 'butterfly', 'cat', 'caterpillar',
    'chimpanzee', 'cockroach', 'cow', 'coyote', 'crab', 'crow', 'deer', 'dog', 'dolphin', 'donkey',
    'dragonfly', 'duck', 'eagle', 'elephant', 'flamingo', 'fly', 'fox', 'goat', 'goldfish', 'goose',
    'gorilla', 'grasshopper', 'hamster', 'hare', 'hedgehog', 'hippopotamus', 'hornbill', 'horse',
    'hummingbird', 'hyena', 'jellyfish', 'kangaroo', 'koala', 'ladybugs', 'leopard', 'lion', 'lizard',
    'lobster', 'mosquito', 'moth', 'mouse', 'octopus', 'okapi', 'orangutan', 'otter', 'owl', 'ox',
    'oyster', 'panda', 'parrot', 'pelecaniformes', 'penguin', 'pig', 'pigeon', 'porcupine', 'possum',
    'raccoon', 'rat', 'reindeer', 'rhinoceros', 'sandpiper', 'seahorse', 'seal', 'shark', 'sheep',
    'snake', 'sparrow', 'squid', 'squirrel', 'starfish', 'swan', 'tiger', 'turkey', 'turtle', 'whale',
    'wolf', 'wombat', 'woodpecker', 'zebra'
]

def classify_animal(image):
    vit_results = vit_classifier(image)
    vit_output = {result['label']: round(result['score'], 3) for result in vit_results}

    clip_results = clip_detector(image, candidate_labels=labels_animals)
    clip_output = {result['label']: round(result['score'], 3) for result in clip_results}

    return {
        "ViT Fine-Tuned": vit_output,
        "CLIP Zero-Shot": clip_output
    }
example_images = [
    ["example_images/lion.jpg"],
    ["example_images/bee.jpg"],
    ["example_images/dog.jpg"]
]

iface = gr.Interface(
    fn=classify_animal,
    inputs=gr.Image(type="filepath"),
    outputs=gr.JSON(),
    title="Animal Classification Comparison",
    description="Upload an image of an animal and compare predictions between a fine-tuned ViT model and a zero-shot CLIP model.",
    examples=example_images
)

iface.launch()