AI-Enabled Waste Segregation System

A two-level AI pipeline for classifying waste as dry or wet, served through a web interface. Combines a custom-trained YOLOv8 object detector with a fine-tuned MobileNetV2 classifier.

---

How It Works

Level 1 — Classifier

A standalone MobileNetV2 binary classifier. Takes a single image and returns DRY or WET with a confidence score. Used directly by the web app's Single mode.

Level 2 — Hybrid Detector

Combines two models in sequence:

1. YOLOv8 (kaustav.pt) detects and localises individual waste objects in the image, returning bounding box coordinates
2. Each detected crop is passed to MobileNetV2 (waste_model_v2.keras) for independent dry/wet classification

YOLO is responsible only for finding and locating objects. MobileNetV2 is responsible for all wet/dry classification decisions — including a full-image fallback when YOLO finds no detections.

Web App

A Flask application with two modes:

• Single — classifies the full image directly via Level 1 (no YOLO involved)
• Multi — runs the full Level 2 hybrid pipeline, returns an annotated image with bounding boxes and a per-object results list

---

Directory Structure

AI-waste-segregation-system/
├── .gitignore
├── README.md
├── requirements.txt
├── download_models.py            ← run this after cloning to get model files
│
├── level1_classifier/
│   ├── model/                    ← place model files here (gitignored)
│   │   ├── waste_model_v2.keras  ← production model, used by web app
│   │   └── dry_wet_model.keras   ← evaluation model, used by analyze_performance.py
│   ├── train.py                  ← training script
│   ├── predict.py                ← single-image inference
│   └── analyze_performance.py   ← evaluation metrics and plots
│
├── level2_detector/
│   ├── kaustav.pt                ← custom YOLOv8 model (gitignored)
│   └── hybrid_detector.py        ← YOLO + MobileNetV2 hybrid pipeline
│
├── tests/
│   ├── test_hybrid.py            ← pytest tests for hybrid_detector
│   └── test_yolo.py              ← pytest tests for YOLO layer
│
└── web_app/
    ├── app.py                    ← Flask application
    ├── templates/
    │   └── index.html            ← UI (Single + Multi modes)
    └── static/
        ├── uploads/              ← user uploaded images (auto-created at runtime)
        └── results/              ← annotated output images (auto-created at runtime)
            └── crops/            ← individual object crops (auto-created at runtime)

---

Model Files

These files are gitignored due to size. Download them using the provided script (see Setup).

File	Location	Used by
waste_model_v2.keras	level1_classifier/model/	Web app (both modes), predict.py, hybrid_detector.py
dry_wet_model.keras	level1_classifier/model/	analyze_performance.py only
kaustav.pt	level2_detector/	hybrid_detector.py

---

Setup

1. Clone the repository

git clone https://github.com/aabirpal/ai-waste-segregation-system.git
cd ai-waste-segregation-system

2. Create a virtual environment

python -m venv venv

# Windows
venv\Scripts\activate

# macOS / Linux
source venv/bin/activate

3. Install dependencies

pip install -r requirements.txt

4. Download model files

python download_models.py

This will download waste_model_v2.keras and kaustav.pt from the GitHub Release and place them in the correct directories automatically. dry_wet_model.keras is optional and only needed if you intend to run analyze_performance.py.

5. Run the web app

cd web_app
python app.py

Open http://127.0.0.1:5000 in your browser.

---

Usage

Web App

Navigate to http://127.0.0.1:5000.

Single mode — Upload an image of a single waste item. MobileNetV2 classifies it directly and returns dry or wet.

Multi mode — Upload an image containing multiple waste items. YOLO detects and crops each object, MobileNetV2 classifies each crop independently, and the UI returns an annotated image with bounding boxes alongside a per-object results list.

Accepted formats: JPG, JPEG, PNG, WEBP, BMP — maximum 16MB.

Command Line

Evaluate model performance:

cd level1_classifier

# Set environment variables for your test dataset and model paths
# Linux / macOS
export TEST_PATH=/path/to/DATASET_BINARY/TEST
export MODEL_PATH=/path/to/model/waste_model_v2.keras

# Windows
set TEST_PATH=C:\path\to\DATASET_BINARY\TEST
set MODEL_PATH=C:\path\to\model\waste_model_v2.keras

python analyze_performance.py

Outputs: classification report, accuracy score, confusion matrix heatmap, ROC curve with AUC. Evaluates waste_model_v2.keras (the production model) by default.

Retrain the Level 1 classifier:

cd level1_classifier
python train.py

Downloads the Kaggle dataset automatically via kagglehub, remaps the original three classes to binary (Organic → Wet, Non-recyclable + Recyclable → Dry), trains a MobileNetV2 model, and saves to level1_classifier/model/dry_wet_model.keras.

---

Running Tests

pytest tests/

Tests are split into two tiers:

• Unit tests — fully mocked, run anywhere without model files present
• Integration tests — marked skipif and skipped automatically when model files are absent (gitignored). Run only when models are available.

---

Dataset

The Level 1 classifier was trained on the Waste Classification Data v2 dataset from Kaggle. The original three-class dataset was remapped to binary:

Original class	Label	Binary
O — Organic	Food and plant waste	Wet
N — Non-recyclable	General waste	Dry
R — Recyclable	Paper, plastic, metal	Dry

The train.py script handles this remapping and dataset download automatically.

---

Contributors

• Aabir Pal
• Ankit Beura
• Kaustav Pramanik
• Satyajeet Sen
• Suryansh Mohanty

---