app.py

# Install required packages
!pip install transformers
!pip install streamlit
!pip install torch
!pip install datasets gradio

import gradio as gr
from transformers import pipeline

# Load a pre-trained image classification pipeline from Hugging Face
model = pipeline("image-classification", model="google/vit-base-patch16-224")

# Define the prediction function
def classify_image(image):
    predictions = model(image)
    return {pred["label"]: pred["score"] for pred in predictions}

# Set up the Gradio interface
interface = gr.Interface(
    fn=classify_image,
    inputs=gr.Image(type="pil"),
    outputs=gr.JSON(),  # Updated to JSON for dictionary-like output
    title="Image Classification App",
    description="Upload an image, and the app will classify it using a Vision Transformer (ViT) model."
)

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