corrected waste categories, added readme and model training notebook and sample images

README.md

@@ -10,8 +10,72 @@ app_file: app.py
 pinned: false
 ---
 
-# Mülltrenner9000
+# Mülltrenner9000 ♻️  
+**Bin there, sorted that.**
 
-Bin there, sorted that.
-An instance segmentation app for waste classification, built with Gradio and YOLOv8.
+An instance segmentation app for waste classification, built with Gradio and YOLOv11.
 
+Mülltrenner9000 is a smart instance segmentation web app that classifies and segments waste into **59 categories**, recommending the correct **German recycling bin**—such as Gelbe Tonne, Restmüll, or Pfand—based on real-time object detection.
+
+The model uses a fine-tuned **YOLOv11 segmentation** model (YOLO11m-seg) and offers an interactive Gradio-based UI for waste classification and bin assignment.
+
+---
+
+🔗 **Try it on Hugging Face Spaces**: [Mülltrenner9000 – Hugging Face Space](https://huggingface.co/spaces/Yas1n/mulltrenner9000)
+
+---
+## 🧠 What It Does
+
+- ✅ **Instance Segmentation** of individual waste items
+- ✅ **Classification** into correct bins:
+  - **Yellow Bin (Gelbe Tonne)** – Plastics & Metals 
+  - **Grey Bin (Restmüll)** – General Waste
+  - **Green Bin (Biotonne)** – Biodegradable Waste 
+  - **Blue Bin (Papiertonne)** – Paper & Cardboard 
+  - **Glascontainer** – Glass Waste
+  - **Hazardous Waste (Sondermüll)** – e.g. Batteries, Chemicals
+  - **Deposit Return (Pfand)** – Refundable items
+- ✅ **Real-time segmentation masks** with bin color overlays and labels
+
+---
+
+## 🗂 Dataset
+
+- **Source:** [TACO: Trash Annotations in Context (Roboflow)](https://universe.roboflow.com/mohamed-traore-2ekkp/taco-trash-annotations-in-context/dataset/15)
+- **Annotations:** Pixel-level masks + class labels
+- **Images:** 3147 for training, 299 for validation
+- **Classes:** 59 object types
+
+---
+
+## 🛠 Model Training
+
+- **Model**: [YOLOv11 (YOLO11m-seg)](https://github.com/ultralytics/ultralytics)
+- **Backbone**: Fine-tuned with heavy data augmentation (RandAugment, CLAHE, etc.)
+- **Epochs**: Trained for 70 epochs
+- **Batch Size**: 16
+- **Optimizer**: SGD
+- **Learning Rate**: Final stage LR tuned to `0.001`
+- **Mixed Precision**: AMP enabled
+
+---
+
+## 🖼 Sample Output
+
+Here is a sample segmentation and classification output of Mülltrenner9000:
+
+![Sample Output](images/sample_image.webp)
+![Sample Output](images/sample_image_1.webp)
+
+
+---
+
+## 🤝 Contribution
+
+We welcome contributions from the community, whether it's in the form of bug reports, improvements, documentation, or new features.
+
+---
+
+## 📄 License
+
+This project is licensed under the **MIT License**. See the `LICENSE` file for more details.

src/model/training_mulltrenner9000.ipynb

@@ -0,0 +1,279 @@
+{
+  "cells": [
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "3uIkSzt52mUL",
+        "outputId": "96aacaf4-2244-4799-e1df-4595c67125c7"
+      },
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "Mounted at /content/drive\n"
+          ]
+        }
+      ],
+      "source": [
+        "from google.colab import drive\n",
+        "drive.mount('/content/drive')"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "0TygyeSi2HZ1"
+      },
+      "outputs": [],
+      "source": [
+        "!pip install torch torchaudio torchvision roboflow ultralytics"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "KlFNecLz2UBJ",
+        "outputId": "c01b326d-7c25-4fd2-e004-70ec98233eed"
+      },
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "loading Roboflow workspace...\n",
+            "loading Roboflow project...\n"
+          ]
+        },
+        {
+          "name": "stderr",
+          "output_type": "stream",
+          "text": [
+            "Downloading Dataset Version Zip in TACO:-Trash-Annotations-in-Context-Dataset-15 to yolov11:: 100%|██████████| 326497/326497 [00:06<00:00, 53159.31it/s]"
+          ]
+        },
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "\n"
+          ]
+        },
+        {
+          "name": "stderr",
+          "output_type": "stream",
+          "text": [
+            "\n",
+            "Extracting Dataset Version Zip to TACO:-Trash-Annotations-in-Context-Dataset-15 in yolov11:: 100%|██████████| 7206/7206 [00:01<00:00, 4606.91it/s]\n"
+          ]
+        }
+      ],
+      "source": [
+        "# Dataset import (https://universe.roboflow.com/mohamed-traore-2ekkp/taco-trash-annotations-in-context/dataset/15)\n",
+        "\n",
+        "from roboflow import Roboflow\n",
+        "rf = Roboflow(api_key=\"Enter your key\") # Enter your roboflow api key here\n",
+        "project = rf.workspace(\"mohamed-traore-2ekkp\").project(\"taco-trash-annotations-in-context\")\n",
+        "version = project.version(15)\n",
+        "dataset = version.download(\"yolov11\")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "colab": {
+          "base_uri": "https://localhost:8080/"
+        },
+        "id": "MtWiembb2fWJ",
+        "outputId": "5886ffd9-d982-4fd1-ee3a-98d922990587"
+      },
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "Creating new Ultralytics Settings v0.0.6 file ✅ \n",
+            "View Ultralytics Settings with 'yolo settings' or at '/root/.config/Ultralytics/settings.json'\n",
+            "Update Settings with 'yolo settings key=value', i.e. 'yolo settings runs_dir=path/to/dir'. For help see https://docs.ultralytics.com/quickstart/#ultralytics-settings.\n"
+          ]
+        }
+      ],
+      "source": [
+        "from ultralytics import YOLO"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "G829VIF32hw1"
+      },
+      "outputs": [],
+      "source": [
+        "#Training the model\n",
+        "\n",
+        "model = YOLO(\"/content/drive/MyDrive/best_heavy_59classes.pt\") # load pretrained model (YOLO11n-seg)\n",
+        "\n",
+        "# Train the model (single gpu)\n",
+        "results = model.train(data=\"/content/TACO:-Trash-Annotations-in-Context-Dataset-15/data.yaml\", epochs=50, imgsz=640)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "iBC5snDI12PQ"
+      },
+      "source": [
+        "### Visualization ###"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "ogkaAT740wVs"
+      },
+      "outputs": [],
+      "source": [
+        "# Waste Bin Mapping (German Mülltrennung System)\n",
+        "bin_map = {\n",
+        "    \"Yellow Bin (Gelbe Tonne)\": [\n",
+        "        \"Aluminium blister pack\", \"Aluminium foil\", \"Carded blister pack\", \"Clear plastic bottle\",\n",
+        "        \"Disposable food container\", \"Disposable plastic cup\", \"Drink Carton\", \"Drink can\",\n",
+        "        \"Foam cup\", \"Meal carton\", \"Metal lid\", \"Metal bottle cap\", \"Other plastic bottle\",\n",
+        "        \"Other plastic container\", \"Other plastic cup\", \"Other plastic wrapper\", \"Plastic bottle cap\",\n",
+        "        \"Plastic film\", \"Plastic glooves\", \"Plastic lid\", \"Plastic straw\", \"Polypropylene bag\",\n",
+        "        \"Pop tab\", \"Single-use carrier bag\", \"Six pack rings\", \"Spread tub\", \"Squeezable tube\",\n",
+        "        \"Tupperware\"\n",
+        "    ],\n",
+        "    \"Grey Bin (Restmüll)\": [\n",
+        "        \"Cigarette\", \"Garbage bag\", \"Shoe\", \"Unlabeled litter\", \"Plastified paper bag\",\n",
+        "        \"Styrofoam piece\", \"Rope & strings\", \"Foam food container\", \"Other plastic\",\n",
+        "        \"Pizza box\", \"Tissues\"\n",
+        "    ],\n",
+        "    \"Green Bin (Biotonne)\": [\n",
+        "        \"Food waste\"\n",
+        "    ],\n",
+        "    \"Blue Bin (Papiertonne)\": [\n",
+        "        \"Egg carton\", \"Normal paper\", \"Other Carton\", \"Paper Bag\", \"Paper cup\", \"Paper straw\",\n",
+        "        \"Pizza box\", \"Toilet tube\", \"Magazine paper\", \"Wrapping paper\", \"Corrugated carton\"\n",
+        "    ],\n",
+        "    \"Glascontainer\": [\n",
+        "        \"Glass bottle\", \"Glass cup\", \"Glass jar\", \"Broken glass\"\n",
+        "    ],\n",
+        "    \"Hazardous Waste (Sondermüll)\": [\n",
+        "        \"Battery\", \"Aerosol\", \"Scrap metal\"\n",
+        "    ],\n",
+        "    \"Deposit Return (Pfand)\": [\n",
+        "        \"Drink can\", \"Clear plastic bottle\", \"Glass bottle\", \"Food Can\"\n",
+        "    ]\n",
+        "}\n",
+        "\n",
+        "# Bin colors in BGR format (for OpenCV)\n",
+        "bin_colors = {\n",
+        "    \"Yellow Bin (Gelbe Tonne)\": (0, 255, 255),       # Yellow\n",
+        "    \"Grey Bin (Restmüll)\": (128, 128, 128),          # Gray\n",
+        "    \"Green Bin (Biotonne)\": (0, 255, 0),             # Green\n",
+        "    \"Blue Bin (Papiertonne)\": (200, 0, 0),           # Blue (bright)\n",
+        "    \"Glascontainer\": (0, 0, 0),                      # Black\n",
+        "    \"Hazardous Waste (Sondermüll)\": (0, 0, 255),     # Red\n",
+        "    \"Deposit Return (Pfand)\": (255, 0, 255)          # Purple/Magenta\n",
+        "}\n",
+        "\n",
+        "class_to_bin = {cls: bin_type for bin_type, cls_list in bin_map.items() for cls in cls_list} #class name gives bin type\n",
+        "\n",
+        "from ultralytics import YOLO\n",
+        "import torch\n",
+        "import cv2\n",
+        "import numpy as np\n",
+        "import matplotlib.pyplot as plt\n",
+        "\n",
+        "# Load model\n",
+        "model = YOLO(\"/content/best_heavy_59classes.pt\") # enter path to your model\n",
+        "\n",
+        "# Load image\n",
+        "image_path = \"/content/IMG-20221229-WA0003_jpg.rf.fce2a8664e953ba379b83eb2e29a32cb.jpg\" # enter image path for testing\n",
+        "image = cv2.imread(image_path)\n",
+        "image = cv2.resize(image, (640, 640))\n",
+        "\n",
+        "# Run prediction\n",
+        "results = model(image_path)\n",
+        "result = results[0]\n",
+        "\n",
+        "# Get class names\n",
+        "class_names = model.names\n",
+        "\n",
+        "# Inverse mapping: class name → bin\n",
+        "class_to_bin = {cls: bin_name for bin_name, class_list in bin_map.items() for cls in class_list}\n",
+        "\n",
+        "# Get results\n",
+        "masks = result.masks.data.cpu().numpy()     # [N, H, W]\n",
+        "boxes = result.boxes.xyxy.cpu().numpy()     # [N, 4]\n",
+        "scores = result.boxes.conf.cpu().numpy()    # [N]\n",
+        "class_ids = result.boxes.cls.cpu().numpy()  # [N]\n",
+        "\n",
+        "overlay = image.copy()\n",
+        "\n",
+        "for i in range(len(class_ids)):\n",
+        "    class_id = int(class_ids[i])\n",
+        "    score = scores[i]\n",
+        "    mask = masks[i]\n",
+        "    box = boxes[i].astype(int)\n",
+        "\n",
+        "    class_name = class_names[class_id]\n",
+        "\n",
+        "    bin_type = next((bin_name for bin_name, items in bin_map.items() if class_name in items), \"Cannot detect trash type\")\n",
+        "    color = bin_colors.get(bin_type, (255, 255, 255))  # fallback to white\n",
+        "\n",
+        "    # Apply colored mask (blended)\n",
+        "    mask_3c = np.stack([mask] * 3, axis=-1)  # [H, W, 3]\n",
+        "    color_array = np.array(color, dtype=np.uint8).reshape(1, 1, 3)\n",
+        "    colored_mask = (mask_3c * color_array).astype(np.uint8)\n",
+        "\n",
+        "    overlay = cv2.addWeighted(overlay, 1.0, colored_mask, 0.5, 0)\n",
+        "\n",
+        "    # Draw bounding box\n",
+        "    cv2.rectangle(overlay, (box[0], box[1]), (box[2], box[3]), color, 2)\n",
+        "\n",
+        "    # Draw label (class name + score + bin)\n",
+        "    label = f\"{class_name} ({int(score * 100)}%)\\n{bin_type}\"\n",
+        "    for j, line in enumerate(label.split(\"\\n\")):\n",
+        "        text_pos = (box[0], box[1] - 10 - 20 * j)\n",
+        "        cv2.putText(overlay, line, text_pos, cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2, cv2.LINE_AA)\n",
+        "\n",
+        "# Show the image\n",
+        "plt.figure(figsize=(14, 10))\n",
+        "plt.imshow(cv2.cvtColor(overlay, cv2.COLOR_BGR2RGB))\n",
+        "plt.axis(\"off\")\n",
+        "plt.title(\"YOLOv11 Mülltrennung Classification Overlay\")\n",
+        "plt.show()"
+      ]
+    }
+  ],
+  "metadata": {
+    "accelerator": "GPU",
+    "colab": {
+      "gpuType": "T4",
+      "provenance": []
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    },
+    "language_info": {
+      "name": "python"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

src/process/process.py

@@ -11,7 +11,7 @@ FILENAME = "best_0.1.0.pt"
 bin_map = {
     "Yellow Bin (Gelbe Tonne)": [
         "Aluminium blister pack", "Aluminium foil", "Carded blister pack", "Clear plastic bottle",
-        "Disposable food container", "Disposable plastic cup", "Drink Carton", "Drink can",
+        "Disposable food container", "Disposable plastic cup", "Drink carton",
         "Foam cup", "Meal carton", "Metal lid", "Metal bottle cap", "Other plastic bottle",
         "Other plastic container", "Other plastic cup", "Other plastic wrapper", "Plastic bottle cap",
         "Plastic film", "Plastic glooves", "Plastic lid", "Plastic straw", "Polypropylene bag",