app.py

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()