import json from functools import lru_cache from urllib.parse import quote_plus import gradio as gr import torch import torch.nn.functional as F from huggingface_hub import hf_hub_download from PIL import Image from torchvision import models, transforms # ── Repos ───────────────────────────────────────────────────────────────────── MOBILENET_REPO = "cpoisson/plantnet300k-mobilenetv3-small" RESNET_REPO = "cpoisson/plantnet300k-resnet18" NUM_CLASSES = 1081 MODEL_CHOICES = { "MobileNetV3-Small (v2) — 10 MB · 3.9M params · +1.56% accuracy ⭐": "mobilenet_v2", "MobileNetV3-Small (v1) — 10 MB · 3.9M params · baseline": "mobilenet", "ResNet18 — 45 MB · 11.7M params · reference model": "resnet18", } # ── Class names ─────────────────────────────────────────────────────────────── _json_path = hf_hub_download(MOBILENET_REPO, "plantnet300K_species_id_2_name.json") with open(_json_path) as f: _id2name = json.load(f) class_ids = sorted(int(k) for k in _id2name) class_names = [_id2name[str(cid)] for cid in class_ids] # ── Preprocessing (ImageNet stats, shared by both models) ───────────────────── transform = 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]), ]) # ── Model loading (lazy, cached) ────────────────────────────────────────────── @lru_cache(maxsize=3) def load_model(key: str) -> torch.nn.Module: if key == "mobilenet_v2": m = models.mobilenet_v3_small(weights=None, num_classes=NUM_CLASSES) path = hf_hub_download(MOBILENET_REPO, "mobilenetv3_small_v2.pth") elif key == "mobilenet": m = models.mobilenet_v3_small(weights=None, num_classes=NUM_CLASSES) path = hf_hub_download(MOBILENET_REPO, "plantnet_mobilenetv3.pth") else: m = models.resnet18(weights=None, num_classes=NUM_CLASSES) path = hf_hub_download(RESNET_REPO, "plantnet_resnet18.pth") m.load_state_dict(torch.load(path, map_location="cpu", weights_only=True)) return m.eval() # ── Inference ───────────────────────────────────────────────────────────────── def classify(image_path: str, model_label: str, top_k: int): if image_path is None: return {}, "" key = MODEL_CHOICES[model_label] model = load_model(key) img = Image.open(image_path).convert("RGB") tensor = transform(img).unsqueeze(0) with torch.no_grad(): probs = F.softmax(model(tensor), dim=1)[0] topk_probs, topk_idx = probs.topk(top_k) predictions = { class_names[i.item()]: float(p) for i, p in zip(topk_idx, topk_probs) } top_name = next(iter(predictions)) search_url = ( "https://www.inaturalist.org/taxa/search?q=" + quote_plus(" ".join(top_name.split()[:2])) ) link_html = ( f'' f'🔍 Search {top_name} on iNaturalist' ) return predictions, link_html # ── About content ───────────────────────────────────────────────────────────── ABOUT = """ ## 🧪 Experiment: small local models for plant identification **Core question** — *How far can a sub-15 MB model go on a real-world, fine-grained botanical dataset?* Can it be useful enough to run entirely offline on a phone or embedded device? This Space presents two fine-tuned models trained on [Pl@ntNet-300K](https://zenodo.org/records/5645731) and evaluated on its held-out test set. > ⚡ **Want to run fully offline in your browser — no server, no internet after first load?** > Try the [offline demo →](https://huggingface.co/spaces/cpoisson/plantnet300k-offline) > Built with ONNX Runtime Web + React, models load from HF Hub and run entirely client-side. --- ### Dataset — Pl@ntNet-300K | | | |---|---| | Source | [Zenodo — DOI:10.5281/zenodo.5645731](https://zenodo.org/records/5645731) | | Paper | Garcin et al., *NeurIPS 2021 Datasets & Benchmarks* | | Images | 306,146 | | Species | **1,081** | | Train split | 245,402 images | | Val split | 29,892 images | | Test split | 31,112 images | | Key challenge | Long-tailed: 80% of species = only 11% of images. High label ambiguity (visually similar species). | --- ### Models — Latest Results | Model | Version | Params | Size | Top-1 (test) | Top-5 (test) | Improvement | |---|---|---|---|---|---|---| | **MobileNetV3-Small** | **v2** ⭐ | 3.9M | **10 MB** | **75.45%** | **93.81%** | +1.56 pp | | MobileNetV3-Small | v1 | 3.9M | 10 MB | 73.89% | 91.86% | baseline | | ResNet18 | v1 | 11.7M | 45 MB | 75.82% | 93.98% | reference | **v2 improvements**: MobileNetV3-Small now surpasses v1 by 1.56 pp through improved training (cosine annealing, class-weighted sampling, TrivialAugment, label smoothing). Still runs at **only 10 MB** — perfect for edge deployment. --- ### Training Details #### v2 (Current) | Parameter | Value | |---|---| | Optimizer | SGD (momentum=0.9, Nesterov=True) | | LR Schedule | Cosine annealing (0.01 → 1e-5) | | Augmentation | TrivialAugment + RandomErasing | | Class Balancing | WeightedRandomSampler | | Label Smoothing | 0.1 | | Epochs | 60 (Phase 1: 5 frozen, Phase 2: 55 full) | | Batch size | 256 | | Best checkpoint | Epoch 59 (val: 75.56%, test: 75.45%) | #### v1 (Previous) | Parameter | Value | |---|---| | Optimizer | Adam, lr = 1e-3 (constant) | | Epochs | 60 | | Batch size | 64 | | Train augmentation | Resize(256) → RandomResizedCrop(224) → HFlip → ColorJitter | | Loss | CrossEntropyLoss | | Checkpoint | Last epoch (no best-val selection) | --- ### Known limitations (v1) - Adam at lr=1e-3 is aggressive for fine-tuning - Weights saved at last epoch, not best checkpoint - Class imbalance not addressed - No LR schedule **v2 addressed all of these**, resulting in **+1.56 pp improvement**. --- ### Replicate ```bash # 1. Download dataset (Zenodo) wget https://zenodo.org/records/5645731/files/plantnet_300K_images.tar.gz tar -xzf plantnet_300K_images.tar.gz # 2. Install pip install torch torchvision # 3. Train (edit DATA_DIR at top of script) python train.py # training script in each model repo ``` Model repos: [cpoisson/plantnet300k-mobilenetv3-small](https://huggingface.co/cpoisson/plantnet300k-mobilenetv3-small) · [cpoisson/plantnet300k-resnet18](https://huggingface.co/cpoisson/plantnet300k-resnet18) """ # ── UI ──────────────────────────────────────────────────────────────────────── sample_images = [[str(p)] for p in sorted(__import__("pathlib").Path("examples").glob("*.jpg"))] with gr.Blocks(title="PlantNet-300K — Small Model Experiment") as demo: gr.Markdown(""" # 🌿 PlantNet-300K — Small Model Experiment Fine-tuned **MobileNetV3-Small (10 MB)** and **ResNet18 (45 MB)** on 1,081 plant species. **v2 now live: 75.45% top-1 accuracy (+1.56 pp improvement)** ⭐ | ⚡ [Run fully offline in your browser →](https://huggingface.co/spaces/cpoisson/plantnet300k-offline) | |---|| """) with gr.Tabs(): # ── Tab 1: Classifier ────────────────────────────────────────────────── with gr.Tab("🔍 Classify"): with gr.Row(): with gr.Column(): image_in = gr.Image( label="Plant photo", type="filepath", sources=["upload", "webcam", "clipboard"], ) model_picker = gr.Radio( choices=list(MODEL_CHOICES.keys()), value=list(MODEL_CHOICES.keys())[0], label="Model", ) top_k = gr.Slider( minimum=1, maximum=10, value=5, step=1, label="Top-K predictions", ) run_btn = gr.Button("Identify 🌱", variant="primary") with gr.Column(): label_out = gr.Label(label="Predicted species", num_top_classes=10) link_out = gr.HTML() gr.Examples( examples=sample_images, inputs=image_in, label="Sample images", ) run_btn.click( fn=classify, inputs=[image_in, model_picker, top_k], outputs=[label_out, link_out], ) image_in.change( fn=classify, inputs=[image_in, model_picker, top_k], outputs=[label_out, link_out], ) # ── Tab 2: About ─────────────────────────────────────────────────────── with gr.Tab("📋 About this experiment"): gr.Markdown(ABOUT) demo.launch()