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)