Spaces:
Sleeping
Sleeping
| import json | |
| from pathlib import Path | |
| import spaces | |
| import gradio as gr | |
| import pandas as pd | |
| import torch | |
| from huggingface_hub import snapshot_download | |
| from model import classify, load_classifier | |
| ROOT = Path(__file__).resolve().parent | |
| deployment = json.loads((ROOT / "deployment.json").read_text()) | |
| model_directory = snapshot_download( | |
| repo_id=deployment["model_id"], | |
| revision=deployment["model_revision"], | |
| allow_patterns=["config.json", "model.safetensors"], | |
| ) | |
| torch.set_num_threads(2) | |
| model, config = load_classifier(model_directory) | |
| model.to("cuda" if torch.cuda.is_available() else "cpu") | |
| def predict(image, top_k): | |
| return { | |
| "model_id": deployment["model_id"], | |
| "model_revision": deployment["model_revision"], | |
| "predictions": classify(model, config, image, int(top_k)), | |
| } | |
| with gr.Blocks(title="Lab 1 — CIFAR-10") as demo: | |
| gr.Markdown("# CIFAR-10 image classifier\nResNet-18 fine-tuned in Google Colab.") | |
| gr.Markdown("Classes: " + ", ".join(config["labels"])) | |
| with gr.Tab("Predict"): | |
| image = gr.Image(type="pil", label="Image", image_mode="RGB") | |
| top_k = gr.Slider(1, len(config["labels"]), value=3, step=1, label="Top K") | |
| button = gr.Button("Predict", variant="primary") | |
| result = gr.JSON(label="Predictions") | |
| button.click(predict, [image, top_k], result, api_name="predict") | |
| with gr.Tab("Metrics"): | |
| summary = json.loads((ROOT / "assets" / "summary.json").read_text()) | |
| gr.Markdown(f"Test accuracy: **{summary['test_accuracy']:.2%}**") | |
| gr.Dataframe(pd.read_csv(ROOT / "assets" / "per_class_metrics.csv"), interactive=False) | |
| gr.Image(str(ROOT / "assets" / "confusion_matrix.png"), label="Confusion matrix", interactive=False) | |
| gr.JSON(summary, label="Training and evaluation details") | |
| if __name__ == "__main__": | |
| demo.queue(default_concurrency_limit=1).launch() | |