Skip to main content

Overview

The Trash Classification AI System is trained to detect and classify three categories of recyclable waste materials. Each category has a unique identifier, color coding, and detection parameters.

Waste Categories

The system recognizes the following waste types:
# From trash_classificator/segmentation/models/trash_model.py
trash_classes: Dict[int, str] = {
    0: 'cardboard and paper',
    1: 'metal',
    2: 'plastic'
}

Class Details

Cardboard & Paper

Class ID: 0
Color: Blue (BGR: 255, 0, 0)
Common items: boxes, newspapers, magazines, paper bags

Metal

Class ID: 1
Color: Yellow (BGR: 255, 255, 0)
Common items: aluminum cans, tin cans, metal containers

Plastic

Class ID: 2
Color: Gray (BGR: 200, 200, 200)
Common items: bottles, containers, packaging

Color Coding System

Each waste category is visualized with a distinct color for easy identification:
# From trash_classificator/drawing/main.py
class MaskDrawer:
    def __init__(self):
        self.color_map: Dict[int, tuple] = {
            0: (255, 0, 0),      # Blue - Cardboard and Paper
            1: (255, 255, 0),    # Yellow - Metal
            2: (200, 200, 200)   # Gray - Plastic
        }
Colors are specified in BGR format (Blue, Green, Red) as per OpenCV convention.

Visual Examples

When trash is detected, the system applies color-coded visualization:
1

Mask Layer

Semi-transparent colored fill (50% opacity) covering the detected object
2

Bounding Box

Solid border around the object in the same category color (2px thickness)
3

Label

Class name displayed above the bounding box

Detection Confidence

The system uses a confidence threshold to filter out low-quality detections:
# From trash_classificator/segmentation/main.py
results = self.trash_segmentation_model.track(
    image, 
    conf=0.55,          # Confidence threshold
    verbose=False, 
    persist=True, 
    imgsz=640,
    stream=True
)
Confidence Threshold: 0.55 (55%)Only detections with confidence scores above 55% are included in the results. This threshold balances between detecting genuine trash objects and avoiding false positives.

Confidence Score Interpretation

Confidence RangeInterpretation
0.90 - 1.00Very high confidence - Clear, unambiguous detection
0.75 - 0.89High confidence - Strong detection with good visibility
0.55 - 0.74Moderate confidence - Valid detection, may have partial occlusion
< 0.55Below threshold - Filtered out as uncertain

Class Mapping and IDs

The class ID system is used throughout the pipeline:

In Segmentation Results

# Class IDs from YOLO results
clss = trash_track.boxes.cls.cpu().tolist()

# Example: [0, 2, 1] means:
# - Object 1: Cardboard and paper (ID: 0)
# - Object 2: Plastic (ID: 2)
# - Object 3: Metal (ID: 1)

In Drawing Operations

# Mask drawing uses class ID to select color
for mask, cls in zip(masks, classes):
    color = self.color_map.get(cls, (255, 255, 255))  # Default: white
    cv2.fillPoly(overlay, mask_polygon, color)

# Bounding box drawing uses class ID to get label
for box, cls in zip(boxes, classes):
    annotator.box_label(box, trash_classes[cls], color=colors(cls, True))
If an unknown class ID is encountered (not 0, 1, or 2), the mask drawer will default to white (255, 255, 255).

Training Data Implications

The three-class system reflects the model’s training data:

Focused Categories

The model specializes in three common recyclable materials, allowing for high accuracy within these categories.

Recyclability Focus

All three categories (paper, metal, plastic) are recyclable materials, making the system ideal for waste sorting applications.

Extending the Class System

To add new waste categories, you would need to:
1

Retrain the Model

Add new labeled training data for the additional category and retrain the YOLO model
2

Update Class Dictionary

Add the new class ID and name to trash_classes in trash_model.py
3

Add Color Mapping

Define a unique color for the new category in MaskDrawer.color_map
4

Test Detection

Verify the new category is detected and visualized correctly
When selecting colors for new categories, choose colors that are visually distinct from existing ones (blue, yellow, gray) for easy differentiation.

Implementation Reference

The complete class configuration is located in:
trash_classificator/
├── segmentation/
│   └── models/
│       └── trash_model.py          # Class definitions
└── drawing/
    └── main.py                      # Color mappings
The trash_classes dictionary is imported and used across multiple modules, ensuring consistent class naming throughout the system.

Build docs developers (and LLMs) love