| import os | |
| import gradio as gr | |
| import timm | |
| from huggingface_hub import login | |
| from torch import no_grad, softmax, topk | |
| MODEL_NAME = os.getenv("MODEL_NAME") | |
| HF_TOKEN = os.getenv("HF_TOKEN") | |
| login(token=HF_TOKEN) | |
| model = timm.create_model(f"hf_hub:{MODEL_NAME}", pretrained=True) | |
| model.eval() | |
| data_cfg = timm.data.resolve_data_config(model.pretrained_cfg) | |
| transform = timm.data.create_transform(**data_cfg) | |
| def classify_image(input): | |
| inp = transform(input) | |
| with no_grad(): | |
| output = model(inp.unsqueeze(0)) | |
| probabilities = softmax(output[0], dim=0) | |
| values, indices = topk(probabilities, 3) | |
| return { | |
| model.pretrained_cfg["label_names"][str(id.item())].title(): prob | |
| for id, prob in zip(indices, values) | |
| } | |
| demo = gr.Interface( | |
| fn=classify_image, | |
| inputs=gr.Image(type="pil", sources=["upload", "clipboard"]), | |
| outputs=gr.Label(num_top_classes=3), | |
| allow_flagging="never", | |
| examples="examples", | |
| ) | |
| demo.queue() | |
| demo.launch(debug=True) | |