README.md

metadata

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

.
├── 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:

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:

chmod 600 ~/.kaggle/kaggle.json

Train Models

Train both baseline and augmented models:

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:

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

Run Local Web App

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:

--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.