app.py

import gradio as gr
from transformers import pipeline

classifier = None

def classify_image(img, top_k):
    if img is None:
        return "Please upload or select an image."

    global classifier
    if classifier is None:
        classifier = pipeline("image-classification", model="google/vit-base-patch16-224")

    try:
        results = classifier(img, top_k=int(top_k))
    except Exception as e:
        return f"Error during classification: {str(e)}"

    output_lines = ["Classification Results:"]
    for index, result in enumerate(results, start=1):
        label = result["label"]
        score = result["score"] * 100
        output_lines.append(f"{index}. {label}: {score:.2f}%")

    return "\n".join(output_lines)


example_images = [
    "animal_images/cat.png",
    "animal_images/frog.png",
    "animal_images/hippo.png",
    "animal_images/jaguar.png",
    "animal_images/sloth.png",
    "animal_images/toucan.png",
    "animal_images/turtle.png",
]

with gr.Blocks(title="Animal Image Classifier") as demo:
    gr.Markdown("# 🐾 Animal Image Classifier")
    gr.Markdown(
        "Upload an animal image or pick one of the examples below. "
        "The model will return the top predictions with confidence scores."
    )

    with gr.Row():
        with gr.Column(scale=2):
            image_input = gr.Image(type="pil", label="Animal Image")
            with gr.Row():
                classify_btn = gr.Button("Classify", variant="primary")
                clear_btn = gr.Button("Clear")
        with gr.Column(scale=1):
            top_k_slider = gr.Slider(
                minimum=1,
                maximum=5,
                step=1,
                value=5,
                label="Top K Predictions",
            )

    output_text = gr.Textbox(
        label="Classification Results",
        lines=8,
        interactive=False,
        container=True,
        elem_id="output-box",
    )

    gr.Examples(
        examples=example_images,
        inputs=image_input,
        label="Example Images",
    )

    demo.css = """
    #output-box textarea {
        background-color: #800000 !important;
        color: white !important;
        font-size: 1.1em !important;
        font-weight: bold !important;
        text-align: center !important;
    }
    """

    classify_btn.click(
        fn=classify_image,
        inputs=[image_input, top_k_slider],
        outputs=output_text,
    )
    clear_btn.click(
        fn=lambda: (None, ""),
        inputs=None,
        outputs=[image_input, output_text],
    )

if __name__ == "__main__":
    demo.launch(share=False)