Overview The Drawing class provides comprehensive visualization for trash detection results. It orchestrates three specialized drawer components to render segmentation masks, bounding boxes with labels, and object tracking trajectories on video frames.
The Drawing class automatically initializes all drawer components and applies them in the optimal order for visual clarity.
Interface Definition DrawingInterfaceAbstract base class defining the contract for drawing systems.
from abc import ABC , abstractmethod import numpy as np from typing import Dict from ultralytics.engine.results import Results class DrawingInterface ( ABC ): @abstractmethod def draw ( self , image : np.ndarray, trash_track : Results, trash_classes : Dict[ int , str ], device ): raise NotImplementedError
Class Definition class Drawing ( DrawingInterface ): def __init__ ( self ): self .mask_drawer: MaskDrawerInterface = MaskDrawer() self .bbox_drawer: BoundingBoxDrawerInterface = BoundingBoxDrawer() self .track_drawer: TrackDrawerInterface = TrackDrawer()
Constructor __init__()Initializes the Drawing system with three specialized drawer components.
Attributes Initialized: Instance of MaskDrawer for rendering segmentation masks with colored overlays
bbox_drawer
BoundingBoxDrawerInterface
Instance of BoundingBoxDrawer for drawing bounding boxes with class labels
Instance of TrackDrawer for visualizing object trajectories and tracking history
Example: from trash_classificator.drawing.main import Drawing # Initialize the drawing system drawer = Drawing()
Methods draw()Renders all detection visualizations on the input image.
def draw ( self , image : np.ndarray, trash_track : Results, trash_classes : Dict[ int , str ], device )
Parameters Input image as a NumPy array in BGR format (OpenCV standard). This image will be modified with visualizations.
YOLO Results object containing detection data:
masks.xy: Polygon coordinates for segmentation masksboxes.xyxy: Bounding box coordinates in [x1, y1, x2, y2] formatboxes.id: Tracking IDs for each detected objectboxes.cls: Class indices for detected trash types Dictionary mapping class indices to human-readable trash type names: { 0 : "plastic" , 1 : "paper" , 2 : "metal" }
PyTorch device object indicating where the model inference was performed. Used for tensor operations during visualization.
Returns The input image with all visualizations applied:
Colored segmentation masks (semi-transparent overlay)
Bounding boxes with class labels
Tracking trajectories showing object movement history
Drawing Pipeline Visualizations are applied in the following order:
Mask Drawing : Semi-transparent colored overlays for segmentationBounding Box Drawing : Boxes with class labels and confidence scoresTrack Drawing : Trajectory lines showing object movement history
Source Code Reference: def draw ( self , image : np.ndarray, trash_track : Results, trash_classes : Dict[ int , str ], device ): masks = trash_track.masks.xy boxes = trash_track.boxes.xyxy.cpu() tracks_ids = trash_track.boxes.id.int().cpu().tolist() clss = trash_track.boxes.cls.cpu().tolist() image = self .mask_drawer.draw(image, masks, clss) image = self .bbox_drawer.draw(image, boxes, trash_classes, clss) image = self .track_drawer.draw(image, tracks_ids, boxes) return image
The draw method modifies the input image in place. Each drawer component receives and returns the progressively modified image.
Drawer Components The Drawing class integrates three specialized drawer components:
MaskDrawer Renders segmentation masks with colored semi-transparent overlays.
Class Definition: class MaskDrawer ( MaskDrawerInterface ): def __init__ ( self ): self .color_map: Dict[ int , tuple ] = { 0 : ( 255 , 0 , 0 ), 1 : ( 255 , 255 , 0 ), 2 : ( 200 , 200 , 200 )}
Color Mapping:
Class 0: Blue (255, 0, 0) in BGR
Class 1: Cyan (255, 255, 0) in BGR
Class 2: Gray (200, 200, 200) in BGR
Implementation: def draw ( self , image : np.ndarray, masks : List[np.ndarray], classes : List[ int ]) -> np.ndarray: overlay = image.copy() for mask, cls in zip (masks, classes): color = self .color_map.get( cls , ( 255 , 255 , 255 )) mask_polygon = np.int32([mask]) cv2.fillPoly(overlay, mask_polygon, color) cv2.polylines(image, mask_polygon, isClosed = True , color = color, thickness = 2 ) alpha = 0.5 return cv2.addWeighted(overlay, alpha, image, 1 - alpha, 0 )
The masks are drawn with 50% transparency (alpha=0.5) for visibility of underlying image. BoundingBoxDrawer Draws bounding boxes with class labels using Ultralytics Annotator.
Show BoundingBoxDrawer Details
Class Definition: class BoundingBoxDrawer ( BoundingBoxDrawerInterface ): def __init__ ( self ): self .thickness: int = 2
Implementation: def draw ( self , image : np.ndarray, boxes : np.ndarray, trash_classes : Dict[ int , str ], classes : List[ int ]) -> np.ndarray: annotator = Annotator(image, self .thickness, trash_classes) for box, cls in zip (boxes, classes): annotator.box_label(box, trash_classes[ cls ], color = colors( cls , True )) return annotator.result()
Features:
Uses Ultralytics color scheme for consistent visualization
Displays class name from trash_classes dictionary
Line thickness of 2 pixels for clear visibility
TrackDrawer Visualizes object trajectories with tracking history.
Class Definition: class TrackDrawer ( TrackDrawerInterface ): def __init__ ( self ): self .track_history = defaultdict( list ) self .thickness: int = 2
Implementation: def draw ( self , image : np.ndarray, tracks_ids : List[ int ], boxes : np.ndarray) -> np.ndarray: for track_id, box in zip (tracks_ids, boxes): track_line = self .track_history[track_id] centroid = ( float ((box[ 0 ] + box[ 2 ]) / 2 ), float ((box[ 1 ] + box[ 3 ]) / 2 )) track_line.append(centroid) if len (track_line) > 50 : track_line.pop( 0 ) for i in range ( 1 , len (track_line)): cv2.line(image, tuple ( map ( int , track_line[i - 1 ])), tuple ( map ( int , track_line[i])), colors(track_id, True ), self .thickness) return image
Features:
Maintains up to 50 historical positions per object
Draws trajectory as connected line segments
Uses unique color per track_id for distinction
Automatically manages memory by limiting history length
Usage Examples Basic Usage import cv2 import numpy as np from trash_classificator.drawing.main import Drawing from trash_classificator.segmentation.main import SegmentationModel # Initialize components segmentation = SegmentationModel() drawer = Drawing() # Load image image = cv2.imread( 'frame.jpg' ) # Run inference results, trash_classes, device = segmentation.inference(image) # Draw visualizations for result in results: if result.boxes.id is not None : visualized_image = drawer.draw(image.copy(), result, trash_classes, device) cv2.imshow( 'Detections' , visualized_image) cv2.waitKey( 0 )
Video Processing with Visualization import cv2 from trash_classificator.drawing.main import Drawing from trash_classificator.segmentation.main import SegmentationModel # Initialize segmentation = SegmentationModel() drawer = Drawing() # Open video cap = cv2.VideoCapture( 'trash_video.mp4' ) while cap.isOpened(): ret, frame = cap.read() if not ret: break # Process frame results, trash_classes, device = segmentation.inference(frame) # Visualize detections for result in results: if result.boxes.id is not None : frame = drawer.draw(frame, result, trash_classes, device) # Display cv2.imshow( 'Trash Detection' , frame) if cv2.waitKey( 1 ) & 0x FF == ord ( 'q' ): break cap.release() cv2.destroyAllWindows()
Custom Visualization Pipeline import cv2 import numpy as np from trash_classificator.drawing.main import ( Drawing, MaskDrawer, BoundingBoxDrawer, TrackDrawer ) from trash_classificator.segmentation.main import SegmentationModel # Initialize with custom drawer configuration drawing = Drawing() segmentation = SegmentationModel() # Access individual drawers for custom control mask_drawer = drawing.mask_drawer bbox_drawer = drawing.bbox_drawer track_drawer = drawing.track_drawer image = cv2.imread( 'scene.jpg' ) results, trash_classes, device = segmentation.inference(image) for result in results: if result.boxes.id is not None : # Extract data masks = result.masks.xy boxes = result.boxes.xyxy.cpu() track_ids = result.boxes.id.int().cpu().tolist() classes = result.boxes.cls.cpu().tolist() # Apply drawers individually for custom control image_with_masks = mask_drawer.draw(image.copy(), masks, classes) image_with_boxes = bbox_drawer.draw(image_with_masks, boxes, trash_classes, classes) final_image = track_drawer.draw(image_with_boxes, track_ids, boxes) cv2.imshow( 'Custom Visualization' , final_image) cv2.waitKey( 0 )
Saving Visualized Frames import cv2 from trash_classificator.drawing.main import Drawing from trash_classificator.segmentation.main import SegmentationModel segmentation = SegmentationModel() drawer = Drawing() cap = cv2.VideoCapture( 'input.mp4' ) fourcc = cv2.VideoWriter_fourcc( * 'mp4v' ) out = cv2.VideoWriter( 'output.mp4' , fourcc, 30.0 , ( 640 , 480 )) frame_count = 0 while cap.isOpened(): ret, frame = cap.read() if not ret: break results, trash_classes, device = segmentation.inference(frame) for result in results: if result.boxes.id is not None : frame = drawer.draw(frame, result, trash_classes, device) out.write(frame) frame_count += 1 print ( f "Processed frame { frame_count } " ) cap.release() out.release() print ( f "Saved { frame_count } frames to output.mp4" )
Integration Details The typical data flow through the Drawing system:
Input : Original image + YOLO Results objectExtraction : Extract masks, boxes, track_ids, and classes from ResultsMask Layer : Apply semi-transparent segmentation masksBox Layer : Add bounding boxes with class labelsTrack Layer : Draw trajectory lines based on tracking historyOutput : Fully annotated image Each layer builds on the previous, creating a comprehensive visualization. Mask Colors (BGR format):
Class 0: Blue (255, 0, 0)
Class 1: Cyan (255, 255, 0)
Class 2: Gray (200, 200, 200)
Unknown: White (255, 255, 255)
Box and Track Colors:
Generated by Ultralytics colors() function
Unique color per class index
Consistent across frames for same class
The color scheme ensures visual distinction between different trash types. Track History:
Stored in defaultdict(list) in TrackDrawer
Limited to 50 positions per track_id
Automatically prunes old positions using FIFO
Persists across frames for continuous tracking
Image Copies:
MaskDrawer creates overlay copy for transparency blending
Original image preserved until final output
Minimal memory overhead due to in-place modifications
Required packages:
numpy: Array operationscv2 (OpenCV): Image drawing and blendingultralytics: Annotator and Results classesshapely>=2.0.0: Polygon operations (checked at import)collections: defaultdict for track history All dependencies are automatically verified at module import.
Sequential Processing : Drawers are applied sequentially, each modifying the imageTensor Operations : Boxes and IDs are moved to CPU before drawingMemory Efficiency : Track history limited to prevent unbounded growthRendering Speed : Drawing operations are CPU-bound and add ~5-10ms per frame For real-time applications, consider:
Using lower resolution images for display
Reducing track history length below 50 positions
Simplifying visualizations by disabling some drawer components
TrashClassificator - Main orchestrator that uses DrawingSegmentationModel - Provides Results data for visualizationMaskDrawerInterface - Interface for mask rendering (drawing/main.py:15)BoundingBoxDrawerInterface - Interface for box rendering (drawing/main.py:37)TrackDrawerInterface - Interface for track rendering (drawing/main.py:54)