refactor: optimize color validation and label drawing with caching and vectorized operations

Author: shoumikchow

bbox_visualizer/core/_utils.py

@@ -2,6 +2,7 @@
 
 from contextlib import contextmanager
 from typing import Generator
+from functools import lru_cache
 
 # Global flag to track warning suppression state
 _warnings_suppressed: bool = False
@@ -63,18 +64,20 @@ def _validate_bbox(bbox: list[int]) -> None:
         )
 
 
+@lru_cache(maxsize=128)
 def _validate_color(color: tuple[int, int, int]) -> None:
-    """Validate BGR color values.
+    """Validate that a color tuple is valid BGR format.
 
     Args:
-        color: BGR color tuple (blue, green, red)
+        color: BGR color tuple to validate
 
     Raises:
-        ValueError: If any color component is outside [0, 255]
-
+        ValueError: If color is not a valid BGR tuple
     """
-    if not all(0 <= c <= 255 for c in color):
-        raise ValueError("Color values must be between 0 and 255")
+    if not isinstance(color, tuple) or len(color) != 3:
+        raise ValueError("Color must be a tuple of 3 integers (BGR)")
+    if not all(isinstance(c, int) and 0 <= c <= 255 for c in color):
+        raise ValueError("Color values must be integers between 0 and 255")
 
 
 def _check_and_modify_bbox(
@@ -101,4 +104,4 @@ def _check_and_modify_bbox(
     bbox = [value if value > 0 else margin for value in bbox]
     bbox[2] = bbox[2] if bbox[2] < img_size[1] else img_size[1] - margin
     bbox[3] = bbox[3] if bbox[3] < img_size[0] else img_size[0] - margin
-    return bbox
+    return bbox

bbox_visualizer/core/labels.py

@@ -3,11 +3,18 @@
 import cv2
 import numpy as np
 from numpy.typing import NDArray
+from functools import lru_cache
 
 from ._utils import _check_and_modify_bbox, _validate_color
 
 font = cv2.FONT_HERSHEY_SIMPLEX
 
+# Cache text size calculations
+@lru_cache(maxsize=128)
+def _get_text_size(label: str, size: float, thickness: int) -> tuple[tuple[int, int], int]:
+    """Get text size with caching for better performance."""
+    return cv2.getTextSize(label, font, size, thickness)
+
 
 def add_label(
     img: NDArray[np.uint8],
@@ -44,7 +51,8 @@ def add_label(
     _validate_color(text_color)
     bbox = _check_and_modify_bbox(bbox, img.shape)
 
-    (text_width, text_height), baseline = cv2.getTextSize(label, font, size, thickness)
+    # Use cached text size calculation
+    (text_width, text_height), baseline = _get_text_size(label, size, thickness)
     padding = 5  # Padding around text
 
     if top and bbox[1] - text_height > text_height:
@@ -54,7 +62,7 @@ def add_label(
 
         # Calculate background rectangle position
         bg_x1 = bbox[0]
-        bg_y1 = bbox[1] - bg_height  # Removed the gap by removing (5 * size)
+        bg_y1 = bbox[1] - bg_height
         bg_x2 = bg_x1 + bg_width
         bg_y2 = bg_y1 + bg_height
 
@@ -127,7 +135,7 @@ def add_multiple_labels(
     text_color: tuple[int, int, int] = (0, 0, 0),
     top: bool = True,
 ) -> NDArray[np.uint8]:
-    """Add multiple labels to their corresponding bounding boxes.
+    """Add multiple labels to their corresponding bounding boxes using optimized operations.
 
     Args:
         img: Input image array
@@ -151,8 +159,25 @@ def add_multiple_labels(
 
     _validate_color(text_bg_color)
     _validate_color(text_color)
+    
+    # Convert bboxes to numpy array for vectorized operations
+    bboxes = np.array(bboxes)
+    
+    # Validate and modify all bboxes at once
+    bboxes = np.array([_check_and_modify_bbox(bbox, img.shape) for bbox in bboxes])
+    
+    # Draw all labels using add_label
+    output = img.copy()
     for label, bbox in zip(labels, bboxes):
-        img = add_label(
-            img, label, bbox, size, thickness, draw_bg, text_bg_color, text_color, top
+        output = add_label(
+            output,
+            label,
+            bbox.tolist(),
+            size,
+            thickness,
+            draw_bg,
+            text_bg_color,
+            text_color,
+            top
         )
-    return img
+    return output

bbox_visualizer/core/rectangle.py

@@ -53,7 +53,7 @@ def draw_multiple_rectangles(
     is_opaque: bool = False,
     alpha: float = 0.5,
 ) -> NDArray[np.uint8]:
-    """Draws multiple rectangles on the image.
+    """Draws multiple rectangles on the image using optimized vectorized operations.
 
     Args:
         img: Input image array
@@ -70,6 +70,22 @@ def draw_multiple_rectangles(
     if not bboxes:
         raise ValueError("List of bounding boxes cannot be empty")
     _validate_color(bbox_color)
+    
+    # Convert bboxes to numpy array for vectorized operations
+    bboxes = np.array(bboxes)
+    
+    # Validate and modify all bboxes at once
+    bboxes = np.array([_check_and_modify_bbox(bbox, img.shape) for bbox in bboxes])
+    
+    # Draw all rectangles using draw_rectangle
+    output = img.copy()
     for bbox in bboxes:
-        img = draw_rectangle(img, bbox, bbox_color, thickness, is_opaque, alpha)
-    return img
+        output = draw_rectangle(
+            output,
+            bbox.tolist(),
+            bbox_color,
+            thickness,
+            is_opaque,
+            alpha
+        )
+    return output