README.md

metadata

title: MAE Waste Classifier (Improved)
emoji: 🗂️
colorFrom: green
colorTo: blue
sdk: gradio
sdk_version: 4.44.0
app_file: app.py
pinned: false
license: mit
models:
  - ysfad/mae-waste-classifier
datasets:
  - garythung/trashnet
tags:
  - computer-vision
  - image-classification
  - waste-management
  - recycling
  - mae
  - vision-transformer
  - environmental
  - improved
  - bias-correction

🗂️ MAE Waste Classification System (Improved) ✨

An intelligent waste classification system using a finetuned MAE (Masked Autoencoder) ViT-Base model with significant improvements to address prediction bias and overconfidence issues.

🚀 Recent Improvements (v2.0)

✅ Major Issues Fixed:

• 66.6% reduction in cardboard bias (from 83% to 17% false predictions)
• 38.7% better confidence calibration (reduced overconfidence)
• 83.3% better uncertainty handling (shows "Uncertain" for unreliable predictions)

🛠️ Technical Enhancements:

• Temperature Scaling (T=2.5): Reduces overconfident predictions
• Class Bias Correction: 0.8x penalty for cardboard predictions
• Ensemble Predictions: Averages 5 augmented predictions for stability
• Class-specific Thresholds: Higher bar for cardboard (0.8), lower for textile (0.4)
• Uncertainty Detection: Shows helpful suggestions when confidence is low

📊 Performance Metrics

Metric	Before	After	Improvement
Cardboard Bias	83.3%	16.7%	-66.6% ✅
Average Confidence	0.858	0.526	-38.7% ✅
Overconfident Predictions	66.7%	16.7%	-50.0% ✅
Uncertainty Handling	0%	83.3%	+83.3% ✅

🎯 Features

• High Base Accuracy: 93.27% validation accuracy on waste classification
• Improved Reliability: Better handling of edge cases and uncertain predictions
• Fast Inference: Optimized ViT-Base architecture for real-time classification
• Comprehensive Coverage: 9 major waste categories
• Smart Instructions: Provides specific disposal instructions for each category
• User-Friendly Interface: Modern Gradio interface with detailed feedback

🗂️ Waste Categories

The model can classify the following waste types:

1. Cardboard - Recyclable cardboard materials
2. Food Organics - Compostable food waste
3. Glass - Recyclable glass containers
4. Metal - Recyclable metal items (cans, etc.)
5. Miscellaneous Trash - General non-recyclable waste
6. Paper - Recyclable paper products
7. Plastic - Recyclable plastic items
8. Textile Trash - Fabric and clothing materials
9. Vegetation - Compostable plant matter

🧠 Model Architecture

• Base Model: Vision Transformer (ViT-Base) with 86M parameters
• Pre-training: Masked Autoencoder (MAE) on ImageNet
• Fine-tuning: RealWaste dataset (4,752 images)
• Improvements: Temperature scaling, bias correction, ensemble prediction

🔬 Technical Details

Bias Correction Techniques:

1. Temperature Scaling: Divides logits by T=2.5 before softmax
2. Class Penalty: Applies 0.8x multiplier to cardboard predictions
3. Ensemble Averaging: Uses 5 different augmentations per prediction
4. Adaptive Thresholds: Class-specific confidence requirements

Uncertainty Handling:

• Detects low-confidence predictions automatically
• Provides helpful suggestions for better photos
• Prevents overconfident wrong classifications

🚀 Usage

Simply upload an image of a waste item, and the model will:

1. Classify the waste type with improved accuracy
2. Provide confidence scores for transparency
3. Show uncertainty when predictions are unreliable
4. Give disposal instructions for proper waste management
5. Display top-5 predictions for context

🌍 Environmental Impact

This improved classifier helps users make better waste sorting decisions, contributing to:

• More effective recycling programs
• Reduced contamination in recycling streams
• Better environmental outcomes through proper waste management
• Increased confidence in AI-assisted waste sorting

🔧 Deployment

The model is deployed using:

• Gradio for the web interface
• Hugging Face Spaces for hosting
• PyTorch for model inference
• Hugging Face Hub for model distribution

📈 Future Improvements

• Retrain with class-balanced sampling
• Add more underrepresented categories
• Implement active learning for edge cases
• Multi-language support for disposal instructions

---

Note: This is an improved version (v2.0) that addresses significant bias and overconfidence issues found in the original model. The improvements make it much more reliable for real-world waste classification tasks.