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 — 10 MB · 3.9M params · edge-deployable ✦": "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=2) def load_model(key: str) -> torch.nn.Module: if 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 | 243,916 images | | Val split | 31,118 images | | Test split | 31,112 images | | Key challenge | Long-tailed: 80% of species = only 11% of images. High label ambiguity (visually similar species). | --- ### Models | Model | Params | Size | Top-1 (test) | Top-5 (test) | |---|---|---|---|---| | **MobileNetV3-Small** | 3.9M | **10 MB** | **73.89%** | **91.86%** | | ResNet18 | 11.7M | 45 MB | 75.82% | 93.98% | MobileNetV3-Small gets the right species in its **top-5 predictions 9 times out of 10** at only **10 MB** — a compelling result for edge deployment. ResNet18 gains +1.9 pp top-1 at 4.5× the size. --- ### Training methodology (v1) | Hyperparameter | Value | |---|---| | Backbone | ImageNet1K_V1 pretrained (torchvision) | | Optimizer | Adam, lr = 1e-3 (constant) | | Epochs | 60 | | Batch size | 64 | | Train augmentation | Resize(256) → RandomResizedCrop(224) → HFlip → ColorJitter | | Val/Test | Resize(256) → CenterCrop(224) | | Loss | CrossEntropyLoss | | Checkpoint | Last epoch (no best-val selection in v1) | **Hardware** — NVIDIA RTX 3070 (8 GB) · Intel i7-8086K · 32 GB RAM · Ubuntu Linux · PyTorch 2.7.0 + CUDA 12.6 --- ### Known limitations of this training run (v1) - Adam at lr=1e-3 is aggressive for fine-tuning — may erode pretrained features - Weights saved at last epoch, not best-val checkpoint - Class imbalance not addressed (no weighted sampling, no label smoothing) - No LR schedule **A v2 training run** incorporating SGD + cosine annealing, two-phase fine-tuning, WeightedRandomSampler, label smoothing and stronger augmentation is currently running. Results will be published here once complete. --- ### 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. Exploring how small a model can be while remaining useful for offline / edge plant identification. This space models can aslo 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()