File size: 1,000 Bytes
a577ef4 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 |
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)
|