ai_playground / app.py
Ravi25's picture
test
66ee980
Raw
History Blame Contribute Delete
1.16 kB
import gradio as gr
from transformers import pipeline
# Initialize the pipeline with the zero-shot image classification model
image_classifier = pipeline(task="zero-shot-image-classification", model="google/siglip-so400m-patch14-384")
# Define the candidate labels (classes) for classification
texts = [
"Mini Dress",
"Midi Dress",
"Maxi Dress",
"Short sleeve",
"Long sleeve",
"Three-Fourth sleeve",
"Puff Sleeve",
"A-line Dress",
"T-shirt Dress",
"Shirt Dress",
"Flowy Dress",
"Halter Neck Dress",
"Cut-out Dress",
]
# Define the prediction function
def predict(input_img):
predictions = image_classifier(input_img, candidate_labels=texts)
# Return the input image and all predictions as a dictionary of label: score
return {p["label"]: p["score"] for p in predictions}
# Set up the Gradio interface
gradio_app = gr.Interface(
predict,
inputs=gr.Image(label="Select Image", sources=['upload', 'webcam'], type="pil"),
outputs=[gr.Label(label="Result")],
title="Image Classification",
)
# Launch the Gradio app
if __name__ == "__main__":
gradio_app.launch(share=True)