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

import torch
import torch.nn as nn
import torchvision.models as models
from torchvision import transforms
from safetensors.torch import load_file

from PIL import Image


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

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


# 2. Rebuild training architecture
class TL_b5(nn.Module):
    def __init__(self, num_classes: int):
        super().__init__()
        self.base_model = models.efficientnet_b5(weights=None)
        self.base_model.classifier = nn.Sequential(
            nn.Linear(2048, 512),
            nn.ReLU(),
            nn.Dropout(0.5),
            nn.Linear(512, num_classes),
        )

    def forward(self, x):
        return self.base_model(x)

model = TL_b5(num_classes=num_classes)


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

state_dict = load_file(checkpoint_path)
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.