from functools import lru_cache from pathlib import Path 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 # ── Config ──────────────────────────────────────────────────────────────────── HF_MODEL_REPO = "cpoisson/trash-optimizer-models" CATEGORY_EMOJI = { "battery": "🔋", "car_battery": "🚗", "cardboard": "📦", "food_organics": "🥦", "glass": "🍾", "light_bulb": "💡", "metal": "🥫", "mirror": "🪞", "miscellaneous_trash":"🗑️", "neon": "🌡️", "paper": "📄", "pharmacy": "💊", "plastic": "♳", "printer_cartridge": "🖨️", "textile_trash": "👕", "tire": "🏎️", "vegetation": "🌿", "wood": "🪵", } 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 ───────────────────────────────────────────────────────────── @lru_cache(maxsize=1) def load_model(): # Resolve latest model folder from pointer file latest_path = hf_hub_download(HF_MODEL_REPO, "latest") with open(latest_path) as f: folder = f.read().strip() # Load class mapping mapping_path = hf_hub_download(HF_MODEL_REPO, f"{folder}/class_mapping.txt") class_to_idx = {} with open(mapping_path) as f: for line in f: name, idx = line.strip().split(":") class_to_idx[name] = int(idx) categories = [k for k, _ in sorted(class_to_idx.items(), key=lambda x: x[1])] # Build model model = models.efficientnet_b0(weights=None) model.classifier[1] = torch.nn.Linear( model.classifier[1].in_features, len(categories) ) # Load weights weights_path = hf_hub_download(HF_MODEL_REPO, f"{folder}/model.pth") model.load_state_dict( torch.load(weights_path, map_location="cpu", weights_only=True) ) model.eval() return model, categories # ── Inference ───────────────────────────────────────────────────────────────── def classify(image, top_k): if image is None: return {} model, categories = load_model() img = Image.fromarray(image).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) return { f"{CATEGORY_EMOJI.get(categories[i.item()], '')} {categories[i.item()]}": float(p) for i, p in zip(topk_idx, topk_probs) } # ── About ───────────────────────────────────────────────────────────────────── ABOUT = """ ## ♻️ Trash Optimizer — Waste Classification Automatic waste classification across **18 categories**, designed to assist in sorting decisions for optimised waste collection routing. --- ### Categories | | | | | |---|---|---|---| | 🔋 Battery | 🚗 Car battery | 📦 Cardboard | 🥦 Food organics | | 🍾 Glass | 💡 Light bulb | 🥫 Metal | 🪞 Mirror | | 🗑️ Misc. trash | 🌡️ Neon | 📄 Paper | 💊 Pharmacy | | ♳ Plastic | 🖨️ Printer cartridge | 👕 Textile | 🏎️ Tire | | 🌿 Vegetation | 🪵 Wood | | | --- ### Model — EfficientNet-B0 (`efficientnet_custom_v2_2`) | | | |---|---| | Architecture | EfficientNet-B0, pretrained ImageNet1K | | Parameters | 4.0M total · 4.0M trainable | | Input size | 224 × 224 RGB | | Classes | 18 waste categories | | **Test accuracy** | **95.07%** | | Inference latency | 5.1 ms / image (CUDA) | This is the latest model as tracked by the [`latest`](https://huggingface.co/cpoisson/trash-optimizer-models/blob/main/latest) pointer in the model repository. --- ### Datasets | Dataset | Source | Categories contributed | |---|---|---| | **RealWaste** | [UCI / Kaggle](https://archive.ics.uci.edu/dataset/908/realwaste) | cardboard, food organics, glass, metal, misc. trash, paper, plastic, textile, vegetation | | **RHWC** | [Kaggle](https://www.kaggle.com/datasets/alistairking/recyclable-and-household-waste-classification) | cardboard, food organics, glass, metal, paper, plastic, textile | | **Custom dataset** | Internal | battery, car battery, light bulb, mirror, neon, pharmacy, printer cartridge, tire, wood | Up to **500 images per category**, min size 150×150px, no augmentation at build time. --- ### Training | Hyperparameter | Value | |---|---| | Optimizer | Adam, lr = 1e-4 | | Weight decay | 1e-4 | | Backbone frozen | All except last 3 blocks | | Epochs | 27 (early stopping) | | Batch size | 32 | | Train augmentation | Resize(256) → RandomCrop(224) → HFlip → ColorJitter | | Val/Test | Resize(256) → CenterCrop(224) | | Loss | CrossEntropyLoss | ### Per-class results (test set) | Category | Precision | Recall | F1 | |---|---|---|---| | battery | 1.000 | 0.993 | 0.996 | | car_battery | 0.996 | 1.000 | 0.998 | | cardboard | 0.906 | 0.917 | 0.912 | | food_organics | 0.944 | 0.937 | 0.940 | | glass | 0.969 | 0.918 | 0.943 | | light_bulb | 1.000 | 0.988 | 0.994 | | metal | 0.883 | 0.886 | 0.884 | | mirror | 0.996 | 1.000 | 0.998 | | miscellaneous_trash | 0.912 | 0.876 | 0.893 | | neon | 0.978 | 0.996 | 0.987 | | paper | 0.900 | 0.863 | 0.881 | | pharmacy | 0.988 | 0.992 | 0.990 | | plastic | 0.800 | 0.868 | 0.833 | | printer_cartridge | 0.967 | 0.996 | 0.981 | | textile_trash | 0.928 | 0.908 | 0.918 | | tire | 0.976 | 0.996 | 0.986 | | vegetation | 0.981 | 0.996 | 0.989 | | wood | 1.000 | 1.000 | 1.000 | --- ### Project Full source code and deployment scripts: [github.com/cpoisson/trash-optimizer](https://github.com/cpoisson/trash-optimizer) Model repository: [cpoisson/trash-optimizer-models](https://huggingface.co/cpoisson/trash-optimizer-models) """ # ── UI ──────────────────────────────────────────────────────────────────────── sample_images = [[str(p)] for p in sorted(Path("examples").glob("*.jpg"))] with gr.Blocks(title="Trash Optimizer") as demo: gr.Markdown(""" # ♻️ Trash Optimizer — Waste Classifier Upload a photo of waste and the model will identify its category across **18 waste types**. Powered by **EfficientNet-B0** (4M params · 95.1% accuracy) trained on RealWaste + custom dataset. """) with gr.Tabs(): with gr.Tab("🔍 Classify"): with gr.Row(): with gr.Column(): image_in = gr.Image( label="Waste photo", sources=["upload", "webcam", "clipboard"], ) top_k = gr.Slider( minimum=1, maximum=10, value=5, step=1, label="Top-K predictions", ) run_btn = gr.Button("Classify ♻️", variant="primary") with gr.Column(): label_out = gr.Label(label="Predicted category", num_top_classes=10) gr.Examples( examples=sample_images, inputs=image_in, label="Sample images", ) run_btn.click(fn=classify, inputs=[image_in, top_k], outputs=label_out) image_in.change(fn=classify, inputs=[image_in, top_k], outputs=label_out) with gr.Tab("📋 About"): gr.Markdown(ABOUT) demo.launch()