Trash-Net

Trash-Net is an image classification vision-language encoder model fine-tuned from google/siglip2-base-patch16-224 for a single-label classification task. It is designed to classify images of waste materials into different categories using the SiglipForImageClassification architecture.

The model categorizes images into six classes:

• Class 0: "cardboard"
• Class 1: "glass"
• Class 2: "metal"
• Class 3: "paper"
• Class 4: "plastic"
• Class 5: "trash"

Classification Report:
              precision    recall  f1-score   support

   cardboard     0.9912    0.9739    0.9825       806
       glass     0.9564    0.9641    0.9602      1002
       metal     0.9523    0.9744    0.9632       820
       paper     0.9520    0.9848    0.9681      1188
     plastic     0.9835    0.9274    0.9546       964
       trash     0.9127    0.9161    0.9144       274

    accuracy                         0.9626      5054
   macro avg     0.9580    0.9568    0.9572      5054
weighted avg     0.9631    0.9626    0.9626      5054

Run with Transformers🤗

!pip install -q transformers torch pillow gradio

import gradio as gr
from transformers import AutoImageProcessor
from transformers import SiglipForImageClassification
from transformers.image_utils import load_image
from PIL import Image
import torch

# Load model and processor
model_name = "prithivMLmods/Trash-Net"
model = SiglipForImageClassification.from_pretrained(model_name)
processor = AutoImageProcessor.from_pretrained(model_name)

def trash_classification(image):
    """Predicts the category of waste material in the image."""
    image = Image.fromarray(image).convert("RGB")
    inputs = processor(images=image, return_tensors="pt")
    
    with torch.no_grad():
        outputs = model(**inputs)
        logits = outputs.logits
        probs = torch.nn.functional.softmax(logits, dim=1).squeeze().tolist()
    
    labels = {
        "0": "cardboard", 
        "1": "glass", 
        "2": "metal", 
        "3": "paper", 
        "4": "plastic", 
        "5": "trash"
    }
    predictions = {labels[str(i)]: round(probs[i], 3) for i in range(len(probs))}
    
    return predictions

# Create Gradio interface
iface = gr.Interface(
    fn=trash_classification,
    inputs=gr.Image(type="numpy"),
    outputs=gr.Label(label="Prediction Scores"),
    title="Trash Classification",
    description="Upload an image to classify the type of waste material."
)

# Launch the app
if __name__ == "__main__":
    iface.launch()

Intended Use:

The Trash-Net model is designed to classify waste materials into different categories. Potential use cases include:

• Waste Management: Assisting in automated waste sorting and recycling.
• Environmental Monitoring: Identifying and categorizing waste in public spaces.
• Educational Purposes: Teaching waste classification and sustainability.
• Smart Cities: Enhancing waste disposal systems through AI-driven classification.