import gradio as gr
from transformers import pipeline
from huggingface_hub import hf_api
from huggingface_hub import DatasetFilter, ModelFilter

api = hf_api.HfApi()
flyswot_models = ModelFilter(author="flyswot")
models = api.list_models(filter=flyswot_models)
models = [model.modelId for model in models]
choices = gr.inputs.Dropdown(choices=models, default='flyswot/convnext-tiny-224_flyswot',label="Model")


def predict(image,model):
    pipe = pipeline("image-classification", model=model)
    predictions = pipe(image, top_k=20)
    return {pred['label']: pred['score'] for pred in predictions}

iface = gr.Interface(
  fn=predict, 
  inputs=[gr.inputs.Image(type='filepath'),choices],
  outputs='label', interpretation='shap')

iface.launch(enable_queue=True)