# Import necessary libraries
from transformers import pipeline
import gradio as gr

# Load a lightweight image classification model
model = pipeline("image-classification", model="facebook/deit-tiny-patch16-224", cache_dir="./model_cache")

# Function to classify an uploaded image
def classify_image(image):
    predictions = model(image)  # Make predictions
    # Format predictions as a dictionary: Label -> Confidence
    return {pred["label"]: round(pred["score"], 4) for pred in predictions}

# Create a Gradio interface for the app
interface = gr.Interface(
    fn=classify_image,          # Function to call
    inputs=gr.Image(type="pil"),  # Input: Image (PIL format)
    outputs=gr.Label(),         # Output: Label with confidence scores
    title="Image Classification App",
    description="Upload an image, and the app will classify it using a vision transformer model."
)

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