import os 
import huggingface_hub as hf_hub
import gradio as gr

client = hf_hub.InferenceClient(token = os.environ['HF_TOKEN'])
client.headers["x-use-cache"] = "0" 

def predict(image, label_list):
    labels = label_list.split(",")
    response = client.zero_shot_image_classification(
        image,
        labels
    )

    result = {v['label'] : v['score'] for v in response}
    return result

app = gr.Interface(
    fn = predict, 
    inputs = [
        gr.Image(type = 'filepath'), 
        gr.Textbox(label = 'Image class list (separated by comma)')
    ], 
    outputs = gr.Label(), 
    title = 'Zero Shot Image Classification',
    description = 'Vinay Kumar Thakur'
)


app.launch()