Skip to main content

​
Overview

The MaskDrawer class implements the MaskDrawerInterface to visualize segmentation masks on images using color-coded overlays with alpha blending. Each trash class is assigned a specific color for easy visual identification.

​
Class Definition

​
MaskDrawer

Implements mask visualization with configurable color mapping for different trash classes. 
from trash_classificator.drawing.main import MaskDrawer
import numpy as np

mask_drawer = MaskDrawer()

​
Attributes

​
color_map
Dict[int, tuple]
Color mapping for trash classes:
  • 0: (255, 0, 0) - Red (e.g., plastic)
  • 1: (255, 255, 0) - Yellow (e.g., paper)
  • 2: (200, 200, 200) - Light gray (e.g., metal)
If a class is not in the map, defaults to white (255, 255, 255).

​
Methods

​
draw

Draws segmentation masks on an image with alpha blending. 
result = mask_drawer.draw(image, masks, classes)

​
Parameters

​
image
np.ndarray
required
Input image in BGR format (OpenCV format). The image will be modified with overlays and contours.
​
masks
List[np.ndarray]
required
List of segmentation masks. Each mask is a numpy array of polygon coordinates representing the mask boundary.
​
classes
List[int]
required
List of class IDs corresponding to each mask. Used to determine the color from color_map.

​
Returns

​
result
np.ndarray
The annotated image with:
  • Filled polygon overlays blended at 50% opacity (alpha = 0.5)
  • Contour lines drawn with 2px thickness
  • Colors based on class IDs from color_map

​
Alpha Blending

The drawer uses OpenCV’s addWeighted function to blend the mask overlay with the original image: 
alpha = 0.5
result = cv2.addWeighted(overlay, alpha, image, 1 - alpha, 0)
This creates a semi-transparent effect where:
  • alpha = 0.5: 50% of the filled mask overlay
  • 1 - alpha = 0.5: 50% of the original image with contours

​
Usage Example

import cv2
import numpy as np
from trash_classificator.drawing.main import MaskDrawer

# Initialize drawer
mask_drawer = MaskDrawer()

# Load image
image = cv2.imread("trash_scene.jpg")

# Example masks (polygon coordinates)
masks = [
    np.array([[100, 100], [200, 100], [200, 200], [100, 200]]),  # Plastic bottle
    np.array([[300, 150], [400, 150], [400, 250], [300, 250]])   # Paper
]

# Class IDs
classes = [0, 1]  # 0=plastic (red), 1=paper (yellow)

# Draw masks
annotated_image = mask_drawer.draw(image, masks, classes)

# Save or display
cv2.imwrite("result.jpg", annotated_image)

​
Interface

​
MaskDrawerInterface

Abstract base class defining the contract for mask drawing implementations. 
from abc import ABC, abstractmethod
import numpy as np
from typing import List

class MaskDrawerInterface(ABC):
    @abstractmethod
    def draw(self, image: np.ndarray, masks: List[np.ndarray], classes: List[int]) -> np.ndarray:
        """draw mask in image"""
        raise NotImplementedError

​
Color Mapping Reference

Class IDColor (BGR)RGBDescription
0(255, 0, 0)BlueFirst trash class
1(255, 255, 0)CyanSecond trash class
2(200, 200, 200)Light GrayThird trash class
Other(255, 255, 255)WhiteDefault fallback

Build docs developers (and LLMs) love

Get started for freeTalk to us