classification / app.py
Mohuu0601's picture
Rename app.pyy to app.py
328ac76 verified
# Install required packages
!pip install transformers
!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.Label(),
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()