refactor: update type hints

Author: shoumikchow

bbox_visualizer/core/_utils.py

@@ -1,6 +1,7 @@
 """Internal utilities for bbox-visualizer."""
 
 from contextlib import contextmanager
+from typing import Generator
 
 # Global flag to track warning suppression state
 _warnings_suppressed: bool = False
@@ -18,7 +19,7 @@ def suppress_warnings(suppress: bool = True) -> None:
 
 
 @contextmanager
-def warnings_suppressed():
+def warnings_suppressed() -> Generator[None, None, None]:
     """Temporarily suppress warnings.
 
     Example:

bbox_visualizer/core/flags.py

@@ -4,6 +4,7 @@
 
 import cv2
 import numpy as np
+from numpy.typing import NDArray
 
 from ._utils import _check_and_modify_bbox, _should_suppress_warning, _validate_color
 from .labels import add_label
@@ -13,15 +14,15 @@
 
 
 def add_T_label(
-    img: np.ndarray,
+    img: NDArray[np.uint8],
     label: str,
     bbox: list[int],
     size: float = 1,
     thickness: int = 2,
     draw_bg: bool = True,
     text_bg_color: tuple[int, int, int] = (255, 255, 255),
     text_color: tuple[int, int, int] = (0, 0, 0),
-) -> np.ndarray:
+) -> NDArray[np.uint8]:
     """Add a T-shaped label with a vertical line connecting to the bounding box.
 
     The label consists of a vertical line extending from the top of the box
@@ -98,7 +99,7 @@ def add_T_label(
 
 
 def draw_flag_with_label(
-    img: np.ndarray,
+    img: NDArray[np.uint8],
     label: str,
     bbox: list[int],
     size: float = 1,
@@ -107,7 +108,7 @@ def draw_flag_with_label(
     line_color: tuple[int, int, int] = (255, 255, 255),
     text_bg_color: tuple[int, int, int] = (255, 255, 255),
     text_color: tuple[int, int, int] = (0, 0, 0),
-) -> np.ndarray:
+) -> NDArray[np.uint8]:
     """Draws a flag-like label with a vertical line and text box.
 
     The flag consists of a vertical line extending from the middle of the box
@@ -183,13 +184,13 @@ def draw_flag_with_label(
 
 
 def add_multiple_T_labels(
-    img: np.ndarray,
+    img: NDArray[np.uint8],
     labels: list[str],
     bboxes: list[list[int]],
     draw_bg: bool = True,
     text_bg_color: tuple[int, int, int] = (255, 255, 255),
     text_color: tuple[int, int, int] = (0, 0, 0),
-) -> np.ndarray:
+) -> NDArray[np.uint8]:
     """Add multiple T-shaped labels to their corresponding bounding boxes.
 
     Args:
@@ -227,14 +228,14 @@ def add_multiple_T_labels(
 
 
 def draw_multiple_flags_with_labels(
-    img: np.ndarray,
+    img: NDArray[np.uint8],
     labels: list[str],
     bboxes: list[list[int]],
     write_label: bool = True,
     line_color: tuple[int, int, int] = (255, 255, 255),
     text_bg_color: tuple[int, int, int] = (255, 255, 255),
     text_color: tuple[int, int, int] = (0, 0, 0),
-) -> np.ndarray:
+) -> NDArray[np.uint8]:
     """Add multiple flag-like labels to their corresponding bounding boxes.
 
     Args:

bbox_visualizer/core/labels.py

@@ -2,14 +2,15 @@
 
 import cv2
 import numpy as np
+from numpy.typing import NDArray
 
 from ._utils import _check_and_modify_bbox, _validate_color
 
 font = cv2.FONT_HERSHEY_SIMPLEX
 
 
 def add_label(
-    img: np.ndarray,
+    img: NDArray[np.uint8],
     label: str,
     bbox: list[int],
     size: float = 1,
@@ -18,7 +19,7 @@ def add_label(
     text_bg_color: tuple[int, int, int] = (255, 255, 255),
     text_color: tuple[int, int, int] = (0, 0, 0),
     top: bool = True,
-) -> np.ndarray:
+) -> NDArray[np.uint8]:
     """Add a label to a bounding box, either above or inside it.
 
     If there isn't enough space above the box, the label is placed inside.
@@ -116,7 +117,7 @@ def add_label(
 
 
 def add_multiple_labels(
-    img: np.ndarray,
+    img: NDArray[np.uint8],
     labels: list[str],
     bboxes: list[list[int]],
     size: float = 1,
@@ -125,7 +126,7 @@ def add_multiple_labels(
     text_bg_color: tuple[int, int, int] = (255, 255, 255),
     text_color: tuple[int, int, int] = (0, 0, 0),
     top: bool = True,
-) -> np.ndarray:
+) -> NDArray[np.uint8]:
     """Add multiple labels to their corresponding bounding boxes.
 
     Args:

bbox_visualizer/core/rectangle.py

@@ -2,18 +2,19 @@
 
 import cv2
 import numpy as np
+from numpy.typing import NDArray
 
 from ._utils import _check_and_modify_bbox, _validate_color
 
 
 def draw_rectangle(
-    img: np.ndarray,
+    img: NDArray[np.uint8],
     bbox: list[int],
     bbox_color: tuple[int, int, int] = (255, 255, 255),
     thickness: int = 3,
     is_opaque: bool = False,
     alpha: float = 0.5,
-) -> np.ndarray:
+) -> NDArray[np.uint8]:
     """Draws a rectangle around an object in the image.
 
     Args:
@@ -45,13 +46,13 @@ def draw_rectangle(
 
 
 def draw_multiple_rectangles(
-    img: np.ndarray,
+    img: NDArray[np.uint8],
     bboxes: list[list[int]],
     bbox_color: tuple[int, int, int] = (255, 255, 255),
     thickness: int = 3,
     is_opaque: bool = False,
     alpha: float = 0.5,
-) -> np.ndarray:
+) -> NDArray[np.uint8]:
     """Draws multiple rectangles on the image.
 
     Args: