File size: 1,887 Bytes
2ea4d01 fe04217 2ea4d01 f203c0b 2ea4d01 f203c0b 2ea4d01 f203c0b 2ea4d01 f203c0b 2ea4d01 fe04217 2ea4d01 fe04217 2ea4d01 f203c0b fe04217 | 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 43 44 45 46 47 48 49 50 51 52 53 | import os
from typing import Dict
from PIL import Image
import torch
import gradio as gr
from transformers import AutoImageProcessor, AutoModelForImageClassification, AutoConfig
# ----- CONFIG -----
MODEL_ID = os.getenv("MODEL_ID", "kritaphatson/thai_snake_image_classifier").strip()
TOPK = int(os.getenv("TOPK", "3"))
torch.set_num_threads(max(1, int(os.getenv("NUM_THREADS", "1"))))
# ----- LOAD -----
# Disable fast processor to avoid the torchvision warning on CPU Spaces
processor = AutoImageProcessor.from_pretrained(MODEL_ID, use_fast=False)
model = AutoModelForImageClassification.from_pretrained(MODEL_ID)
model.eval().to("cpu")
cfg = AutoConfig.from_pretrained(MODEL_ID)
id2label: Dict[int, str] = {int(k): v for k, v in getattr(cfg, "id2label", {}).items()} or {
i: f"class_{i}" for i in range(model.config.num_labels)
}
# ----- PREDICT -----
@torch.inference_mode()
def predict(img: Image.Image) -> Dict[str, float]:
if img.mode != "RGB":
img = img.convert("RGB")
inputs = processor(images=img, return_tensors="pt")
probs = model(**inputs).logits.softmax(dim=-1)[0]
values, indices = torch.topk(probs, k=min(TOPK, probs.numel()))
return {id2label[int(i)]: float(v) for v, i in zip(values, indices)}
# ----- UI -----
DESCRIPTION = """
Upload a snake image. The model returns the top species predictions.
Model: **kritaphatson/thai_snake_image_classifier**.
"""
with gr.Blocks(title="Thai Snake Classifier") as demo:
gr.Markdown("# Thai Snake Classifier")
gr.Markdown(DESCRIPTION)
with gr.Row():
inp = gr.Image(type="pil", label="Upload image", sources=["upload", "clipboard"])
out = gr.Label(label=f"Top {TOPK} predictions")
gr.Button("Classify").click(fn=predict, inputs=inp, outputs=out)
if __name__ == "__main__":
# Gradio 4: queue() takes no concurrency args
demo.queue(max_size=8).launch()
|