4w4kt/waste-classifier-b5

Waste Classifier (EfficientNet-B5)

This model is a waste classifier based on the EfficientNet-B5 architecture, implemented in pure PyTorch (no Hugging Face transformers dependency). It was trained using transfer learning and fine-tuning to categorize waste images into distinct material classes.

Model Details

• Base Architecture: EfficientNet-B5 (torchvision.models.efficientnet_b5)
• Framework: PyTorch
• Task: Multi-class Waste Image Classification (image-classification)
• Input: RGB images (resized and normalized according to standard ImageNet statistics)
• Output: Class probabilities and predicted category according to labels.json

Repository Structure

• model.pth: PyTorch state dictionary (state_dict) containing the best model weights obtained during training.
• labels.json: Index-to-class mapping for target waste categories.

Usage and Inference

To load the model and run inference in a standalone Python script:

import json
from PIL import Image
import torch
import torch.nn as nn
import torchvision.models as models
from torchvision import transforms

# 1. Load label mapping
with open("labels.json", "r") as f:
    label_data = json.load(f)

id2label = label_data["id2label"]
num_classes = len(id2label)

# 2. Rebuild base architecture
model = models.efficientnet_b5(weights=None)
in_features = model.classifier[1].in_features
model.classifier[1] = nn.Linear(in_features, num_classes)

# 3. Load weights dynamically (CPU / CUDA)
checkpoint_path = "model.pth"
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

state_dict = torch.load(checkpoint_path, map_location=device, weights_only=True)
model.load_state_dict(state_dict)
model.to(device)
model.eval()

# 4. Preprocessing pipeline
inference_transforms = transforms.Compose([
    transforms.Resize((456, 456)),
    transforms.ToTensor(),
    transforms.Normalize(
        mean=[0.485, 0.456, 0.406],
        std=[0.229, 0.224, 0.225]
    ),
])

# 5. Run inference
image = Image.open("path/to/image.jpg").convert("RGB")
input_tensor = inference_transforms(image).unsqueeze(0).to(device)

with torch.no_grad():
    outputs = model(input_tensor)
    probabilities = torch.nn.functional.softmax(outputs[0], dim=0)
    top_pred = torch.argmax(probabilities).item()

predicted_label = id2label[str(top_pred)]
confidence = probabilities[top_pred].item() * 100

print(f"Prediction: {predicted_label} ({confidence:.2f}%)")

Training Hyperparameters

• Epochs: 10
• Batch Size: 32
• Learning Rate: 0.0001
• Optimizer: Adam
• Loss Function: CrossEntropyLoss
• Input Resolution: 456x456 px

Limitations and Intended Use

Designed primarily for single-item waste classification where the object is centered in the frame.

High background clutter, poor lighting, or multiple overlapping objects may affect classification accuracy.