import gradio as gr
import os

from transformers import pipeline

image_classifier = pipeline(task="zero-shot-image-classification", model="google/siglip-so400m-patch14-384")
labels = ['street sweeping', 'litter pickup', 'pothole repair', 'parking enforcement']

def image_mod(image):
    outputs = image_classifier(image, candidate_labels=labels)
    result = {dic["label"]: dic["score"] for dic in outputs}
    return result

app = gr.Interface(
    image_mod,
    gr.Image(type="pil"),
    gr.Label(),
    examples=[
        os.path.join(os.path.dirname(__file__), "images/garbage-pickup.jpg"),
        os.path.join(os.path.dirname(__file__), "images/litter.jpg"),
        os.path.join(os.path.dirname(__file__), "images/wrong-park.jpg"),
        os.path.join(os.path.dirname(__file__), "images/pothole.jpg"),
    ],
)

if __name__ == "__main__":
    app.launch()