---
title: Robust Trash Classifier
emoji: 🗑️
colorFrom: blue
colorTo: green
sdk: gradio
sdk_version: 5.34.2
python_version: "3.11"
app_file: app.py
pinned: false
---
# Robust Trash Classification with Transfer Learning

This project prototypes a computer vision system for classifying waste into six TrashNet categories: `cardboard`, `glass`, `metal`, `paper`, `plastic`, and `trash`.

The main goal is not only clean-image accuracy. The goal is to test whether domain-relevant augmentation helps the model remain reliable under real-world noise such as blur, poor lighting, partial occlusion, and camera variation.

## Project Summary

- **Problem:** Trash sorting models can fail when images are blurry, poorly lit, partially occluded, or taken from unusual angles.
- **Dataset:** TrashNet from Kaggle: `feyzazkefe/trashnet`.
- **Model:** Pretrained ResNet-18.
- **Transfer learning approach:** Replace the final classifier layer with a six-class head, freeze most of the backbone, and fine-tune `layer4` plus the classifier head.
- **Augmentations:** Horizontal flip, small rotation, mild color jitter, small affine transform, and light random erasing.
- **Evaluation:** Compare baseline vs. augmented models on both clean and corrupted test images.

## Repository Structure

```text
.
├── train_eval_final.py          # Train baseline and augmented models
├── app.py                       # Gradio web app for deployment
├── visualize_predictions.py     # Save clean/corrupted prediction examples
├── requirements.txt
├── README.md
└── assets/                      # Confusion matrices and presentation images
```

## Setup

Create and activate a virtual environment:

```bash
python -m venv .venv
source .venv/bin/activate
pip install -r requirements.txt
```

Set up Kaggle access:

1. Create a Kaggle API token from your Kaggle account settings.
2. Put `kaggle.json` in `~/.kaggle/kaggle.json`.
3. Run:

```bash
chmod 600 ~/.kaggle/kaggle.json
```

## Train Models

Train both baseline and augmented models:

```bash
python train_eval_final.py --run-mode both --epochs 10
```

Outputs are saved to `out/`, including:

- `best_baseline.pth`
- `best_augmented.pth`
- `results_baseline.json`
- `results_augmented.json`
- `results_comparison.json`
- confusion matrix images
- training curve image

## Generate Example Predictions

After training, generate a visual comparison of clean vs. corrupted predictions:

```bash
python visualize_predictions.py \
  --data-dir ./data \
  --model-path ./out/best_augmented.pth \
  --output ./out/example_predictions.png
```

## Run Local Web App

```bash
python app.py --model-path ./out/best_augmented.pth
```

Then open the local Gradio link in your browser and upload a waste image.

## Deploy to Hugging Face Spaces

1. Create a new Hugging Face Space.
2. Choose **Gradio** as the SDK.
3. Upload these files:
   - `app.py`
   - `requirements.txt`
   - `out/best_augmented.pth`
4. In the Space, set the model path to match the checkpoint location. If you upload the model to the root folder, change the default in `app.py` to:

```python
--model-path best_augmented.pth
```

## Current Findings

The baseline model achieved slightly higher clean accuracy, while the augmented model became more competitive under corrupted conditions and improved some difficult categories such as plastic. The confusion matrices show that glass remains the hardest class because it is often confused with plastic or metal.

This is a useful result for the hackathon: augmentation does not simply maximize clean accuracy; it changes the model’s robustness behavior under realistic noise.

## Suggested Pitch Narrative

1. **Problem:** Waste images in real deployments are messy, blurry, and inconsistent.
2. **Approach:** Use transfer learning from ImageNet with ResNet-18.
3. **Augmentation:** Simulate realistic deployment noise: lighting changes, slight rotations, blur, occlusion, and camera artifacts.
4. **Result:** Compare baseline and augmented models on clean and corrupted test sets.
5. **Takeaway:** Robustness requires evaluating more than clean accuracy.