Spaces:
Sleeping
Sleeping
| import os | |
| import json | |
| import traceback | |
| import spaces | |
| import gradio as gr | |
| import torch | |
| from torchvision import transforms | |
| from huggingface_hub import hf_hub_download | |
| os.environ.setdefault("OMP_NUM_THREADS", "1") | |
| torch.set_num_threads(1) | |
| REPO_ID = "Pablogps/kedar-15-birds" | |
| # Download model and labels | |
| model_path = hf_hub_download( | |
| repo_id=REPO_ID, | |
| filename="model.pt" | |
| ) | |
| labels_path = hf_hub_download( | |
| repo_id=REPO_ID, | |
| filename="labels.json" | |
| ) | |
| with open(labels_path, "r") as f: | |
| labels = json.load(f) | |
| # Load TorchScript model | |
| model = torch.jit.load(model_path, map_location="cpu") | |
| model.eval() | |
| model.to("cuda") | |
| # Inference preprocessing | |
| preprocess = transforms.Compose([ | |
| transforms.Resize(256), | |
| transforms.CenterCrop(224), | |
| transforms.ToTensor(), | |
| transforms.Normalize( | |
| mean=[0.485, 0.456, 0.406], | |
| std=[0.229, 0.224, 0.225], | |
| ), | |
| ]) | |
| def predict(img): | |
| try: | |
| if img is None: | |
| return {} | |
| img = img.convert("RGB") | |
| x = preprocess(img).unsqueeze(0).to("cuda") | |
| with torch.inference_mode(): | |
| logits = model(x) | |
| probs = torch.softmax(logits, dim=1)[0] | |
| return { | |
| labels[i]: float(probs[i]) | |
| for i in range(len(labels)) | |
| } | |
| except Exception as e: | |
| tb = traceback.format_exc() | |
| print(tb, flush=True) | |
| raise gr.Error(f"{type(e).__name__}: {e}") | |
| title = "Bird classifier" | |
| description = ( | |
| "A small model that classifies birds, trained quickly " | |
| "with a subset of kedar 200 using 15 classes" | |
| ) | |
| examples = [ | |
| "examples_american_crow.jpg", | |
| "examples_barn_swallow.jpg", | |
| "examples_pied_kingfisher.jpg", | |
| ] | |
| demo = gr.Interface( | |
| fn=predict, | |
| inputs=gr.Image(type="pil"), | |
| outputs=gr.Label(num_top_classes=3), | |
| title=title, | |
| description=description, | |
| examples=examples, | |
| cache_examples=False, | |
| ) | |
| demo.queue(max_size=8).launch( | |
| show_error=True, | |
| debug=True | |
| ) |