πŸ—‘οΈ Smart Waste Classification Model

A fine-tuned MobileNetV2 model for classifying waste items into 6 categories using computer vision.

Model Performance

  • Validation Accuracy: 97.46%
  • Framework: PyTorch
  • Architecture: MobileNetV2

Classes

Class Description Color
πŸ”΅ plastic Bottles, bags, containers Blue
πŸ“„ paper Newspapers, cardboard, magazines Brown
πŸ”˜ metal Cans, foil, batteries Gray
πŸ’š glass Bottles, jars Green
🟒 organic Food waste, plant matter Dark Green
⚫ non-recyclable Mixed/contaminated waste Black

Quick Usage

import torch
from torchvision import models, transforms
from PIL import Image
from huggingface_hub import hf_hub_download

# Download model
model_path = hf_hub_download(repo_id="karthikeya09/smart_image_recognation", filename="best_model.pth")

# Load model
model = models.mobilenet_v2(weights=None)
model.classifier = torch.nn.Sequential(
    torch.nn.Dropout(p=0.2),
    torch.nn.Linear(1280, 6)
)
checkpoint = torch.load(model_path, map_location='cpu')
model.load_state_dict(checkpoint['model_state_dict'])
model.eval()

# Define transforms
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

# Predict
classes = ['glass', 'metal', 'non-recyclable', 'organic', 'paper', 'plastic']
image = Image.open('your_image.jpg').convert('RGB')
input_tensor = transform(image).unsqueeze(0)

with torch.no_grad():
    outputs = model(input_tensor)
    probs = torch.nn.functional.softmax(outputs, dim=1)
    confidence, predicted = torch.max(probs, 1)

print(f'Predicted: {classes[predicted.item()]} ({confidence.item()*100:.1f}%)')

Training Details

  • Dataset: ~21,000 waste images
  • Training Split: 70% train, 15% val, 15% test
  • Optimizer: Adam (lr=0.001)
  • Class Weights: Used to handle class imbalance
  • Data Augmentation: Random crop, flip, rotation, color jitter
  • Input Size: 224x224 RGB

Dataset Distribution

Category Images
Organic 6,620
Glass 4,022
Paper 3,882
Metal 3,428
Plastic 1,870
Non-recyclable 1,394

Model Architecture

MobileNetV2 (pretrained on ImageNet)
└── classifier
    β”œβ”€β”€ Dropout(p=0.2)
    └── Linear(1280, 6)

License

MIT License

Author

K Karthikeya Gupta

Downloads last month
8
Inference Providers NEW
This model isn't deployed by any Inference Provider. πŸ™‹ Ask for provider support

Space using karthikeya09/smart_image_recognation 1