From 6c2c3f8862e603ec2da55e2ff94ecfdcfa480a2a Mon Sep 17 00:00:00 2001 From: jirka <6035284+Borda@users.noreply.github.com> Date: Mon, 2 Feb 2026 09:50:08 +0100 Subject: [PATCH] Convert various modules' type annotations + docstring examples to align with `numpy.typing` & fix linting --- supervision/dataset/utils.py | 3 +- supervision/detection/tools/csv_sink.py | 4 +- supervision/detection/utils/iou_and_nms.py | 8 +++- supervision/draw/base.py | 3 +- supervision/draw/color.py | 6 ++- supervision/draw/utils.py | 52 +++++++++++++--------- supervision/geometry/utils.py | 3 +- supervision/utils/image.py | 8 ++-- 8 files changed, 55 insertions(+), 32 deletions(-) diff --git a/supervision/dataset/utils.py b/supervision/dataset/utils.py index 02fb3b92..8cf4682f 100644 --- a/supervision/dataset/utils.py +++ b/supervision/dataset/utils.py @@ -234,7 +234,8 @@ def mask_to_rle(mask: npt.NDArray[np.bool_]) -> list[int]: >>> [int(x) for x in rle] [5, 2, 2, 2, 5] - ![mask_to_rle](https://media.roboflow.com/supervision-docs/mask-to-rle.png){ align=center width="800" } + ![mask_to_rle](https://media.roboflow.com/supervision-docs/ + mask-to-rle.png){ align=center width="800" } """ assert mask.ndim == 2, "Input mask must be 2D" assert mask.size != 0, "Input mask cannot be empty" diff --git a/supervision/detection/tools/csv_sink.py b/supervision/detection/tools/csv_sink.py index 8431486c..a81899f5 100644 --- a/supervision/detection/tools/csv_sink.py +++ b/supervision/detection/tools/csv_sink.py @@ -45,7 +45,9 @@ class CSVSink: ... class_id=np.array([0, 1]) ... ) >>> # Use temporary file - >>> temp_file = tempfile.NamedTemporaryFile(mode='w', suffix='.csv', delete=False) + >>> temp_file = tempfile.NamedTemporaryFile( + ... mode='w', suffix='.csv', delete=False + ... ) >>> temp_file.close() >>> csv_sink = sv.CSVSink(temp_file.name) >>> with csv_sink as sink: diff --git a/supervision/detection/utils/iou_and_nms.py b/supervision/detection/utils/iou_and_nms.py index f67fd850..6e1b4ded 100644 --- a/supervision/detection/utils/iou_and_nms.py +++ b/supervision/detection/utils/iou_and_nms.py @@ -194,10 +194,14 @@ def box_iou_batch( ... [150, 150, 250, 250], ... [320, 320, 420, 420] ... ]) - >>> sv.box_iou_batch(boxes_true, boxes_detection, overlap_metric=sv.OverlapMetric.IOU) + >>> sv.box_iou_batch( + ... boxes_true, boxes_detection, overlap_metric=sv.OverlapMetric.IOU + ... ) array([[0.14285..., 0. ], [0. , 0.47058...]], dtype=float32) - >>> sv.box_iou_batch(boxes_true, boxes_detection, overlap_metric=sv.OverlapMetric.IOS) + >>> sv.box_iou_batch( + ... boxes_true, boxes_detection, overlap_metric=sv.OverlapMetric.IOS + ... ) array([[0.25, 0. ], [0. , 0.64]], dtype=float32) """ diff --git a/supervision/draw/base.py b/supervision/draw/base.py index e27c1d3c..0f3d50c6 100644 --- a/supervision/draw/base.py +++ b/supervision/draw/base.py @@ -1,9 +1,10 @@ from typing import TypeVar import numpy as np +import numpy.typing as npt from PIL import Image -ImageType = TypeVar("ImageType", np.ndarray, Image.Image) +ImageType = TypeVar("ImageType", npt.NDArray[np.uint8], Image.Image) """ An image of type `np.ndarray` or `PIL.Image.Image`. diff --git a/supervision/draw/color.py b/supervision/draw/color.py index 2e82a07b..f3d13510 100644 --- a/supervision/draw/color.py +++ b/supervision/draw/color.py @@ -307,7 +307,8 @@ class ColorPalette: Example: >>> import supervision as sv - >>> sv.ColorPalette.from_hex(['#ff0000', '#00ff00', '#0000ff']) # doctest: +ELLIPSIS + >>> colors = ['#ff0000', '#00ff00', '#0000ff'] + >>> sv.ColorPalette.from_hex(colors) # doctest: +ELLIPSIS ColorPalette(colors=[Color(r=255, g=0, b=0), Color(r=0, g=255, b=0), ...]) """ colors = [Color.from_hex(color_hex) for color_hex in color_hex_list] @@ -356,7 +357,8 @@ class ColorPalette: Example: >>> import supervision as sv - >>> color_palette = sv.ColorPalette.from_hex(['#ff0000', '#00ff00', '#0000ff']) + >>> colors = ['#ff0000', '#00ff00', '#0000ff'] + >>> color_palette = sv.ColorPalette.from_hex(colors) >>> color_palette.by_idx(1) Color(r=0, g=255, b=0) """ diff --git a/supervision/draw/utils.py b/supervision/draw/utils.py index f9f72a8e..979e1cb6 100644 --- a/supervision/draw/utils.py +++ b/supervision/draw/utils.py @@ -5,18 +5,19 @@ from typing import cast import cv2 import numpy as np +import numpy.typing as npt from supervision.draw.color import Color from supervision.geometry.core import Point, Rect def draw_line( - scene: np.ndarray, + scene: npt.NDArray[np.uint8], start: Point, end: Point, color: Color = Color.ROBOFLOW, thickness: int = 2, -) -> np.ndarray: +) -> npt.NDArray[np.uint8]: """ Draws a line on a given scene. @@ -41,8 +42,11 @@ def draw_line( def draw_rectangle( - scene: np.ndarray, rect: Rect, color: Color = Color.ROBOFLOW, thickness: int = 2 -) -> np.ndarray: + scene: npt.NDArray[np.uint8], + rect: Rect, + color: Color = Color.ROBOFLOW, + thickness: int = 2, +) -> npt.NDArray[np.uint8]: """ Draws a rectangle on an image. @@ -66,8 +70,11 @@ def draw_rectangle( def draw_filled_rectangle( - scene: np.ndarray, rect: Rect, color: Color = Color.ROBOFLOW, opacity: float = 1 -) -> np.ndarray: + scene: npt.NDArray[np.uint8], + rect: Rect, + color: Color = Color.ROBOFLOW, + opacity: float = 1, +) -> npt.NDArray[np.uint8]: """ Draws a filled rectangle on an image. @@ -105,11 +112,11 @@ def draw_filled_rectangle( def draw_rounded_rectangle( - scene: np.ndarray, + scene: npt.NDArray[np.uint8], rect: Rect, color: Color, border_radius: int, -) -> np.ndarray: +) -> npt.NDArray[np.uint8]: """ Draws a rounded rectangle on an image. @@ -157,11 +164,11 @@ def draw_rounded_rectangle( def draw_polygon( - scene: np.ndarray, - polygon: np.ndarray, + scene: npt.NDArray[np.uint8], + polygon: npt.NDArray[np.int_], color: Color = Color.ROBOFLOW, thickness: int = 2, -) -> np.ndarray: +) -> npt.NDArray[np.uint8]: """Draw a polygon on a scene. Parameters: @@ -180,11 +187,11 @@ def draw_polygon( def draw_filled_polygon( - scene: np.ndarray, - polygon: np.ndarray, + scene: npt.NDArray[np.uint8], + polygon: npt.NDArray[np.int_], color: Color = Color.ROBOFLOW, opacity: float = 1, -) -> np.ndarray: +) -> npt.NDArray[np.uint8]: """Draw a filled polygon on a scene. Parameters: @@ -209,7 +216,7 @@ def draw_filled_polygon( def draw_text( - scene: np.ndarray, + scene: npt.NDArray[np.uint8], text: str, text_anchor: Point, text_color: Color = Color.BLACK, @@ -218,7 +225,7 @@ def draw_text( text_padding: int = 10, text_font: int = cv2.FONT_HERSHEY_SIMPLEX, background_color: Color | None = None, -) -> np.ndarray: +) -> npt.NDArray[np.uint8]: """ Draw text with background on a scene. @@ -246,7 +253,9 @@ def draw_text( >>> from supervision.draw.utils import draw_text >>> scene = np.zeros((100, 100, 3), dtype=np.uint8) >>> text_anchor = Point(x=50, y=50) - >>> scene = draw_text(scene=scene, text="Hello, world!", text_anchor=text_anchor) + >>> scene = draw_text( + ... scene=scene, text="Hello, world!", text_anchor=text_anchor + ... ) >>> scene.shape (100, 100, 3) """ @@ -285,8 +294,11 @@ def draw_text( def draw_image( - scene: np.ndarray, image: str | np.ndarray, opacity: float, rect: Rect -) -> np.ndarray: + scene: npt.NDArray[np.uint8], + image: str | npt.NDArray[np.uint8], + opacity: float, + rect: Rect, +) -> npt.NDArray[np.uint8]: """ Draws an image onto a given scene with specified opacity and dimensions. @@ -329,7 +341,7 @@ def draw_image( # Resize and isolate alpha channel image = cv2.resize(image, (rect_width, rect_height)) - image = cast(np.ndarray, image) + image = cast(npt.NDArray[np.uint8], image) alpha_channel = ( image[:, :, 3] if image.shape[2] == 4 diff --git a/supervision/geometry/utils.py b/supervision/geometry/utils.py index 81e1e5d3..1234dd2e 100644 --- a/supervision/geometry/utils.py +++ b/supervision/geometry/utils.py @@ -1,9 +1,10 @@ import numpy as np +import numpy.typing as npt from supervision.geometry.core import Point -def get_polygon_center(polygon: np.ndarray) -> Point: +def get_polygon_center(polygon: npt.NDArray[np.float64]) -> Point: """ Calculate the center of a polygon. The center is calculated as the center of the solid figure formed by the points of the polygon diff --git a/supervision/utils/image.py b/supervision/utils/image.py index cf4e3808..d8417225 100644 --- a/supervision/utils/image.py +++ b/supervision/utils/image.py @@ -113,7 +113,7 @@ def scale_image(image: ImageType, scale_factor: float) -> ImageType: width_old, height_old = image.shape[1], image.shape[0] width_new = int(width_old * scale_factor) height_new = int(height_old * scale_factor) - return cv2.resize(image, (width_new, height_new), interpolation=cv2.INTER_LINEAR) # type: ignore + return cv2.resize(image, (width_new, height_new), interpolation=cv2.INTER_LINEAR) @ensure_cv2_image_for_standalone_function @@ -170,7 +170,7 @@ def resize_image( else: width_new, height_new = resolution_wh - return cv2.resize(image, (width_new, height_new), interpolation=cv2.INTER_LINEAR) # type: ignore + return cv2.resize(image, (width_new, height_new), interpolation=cv2.INTER_LINEAR) @ensure_cv2_image_for_standalone_function @@ -233,7 +233,7 @@ def letterbox_image( image[:, :padding_left, 3] = 0 image[:, width_new - padding_right :, 3] = 0 - return image_with_borders # type: ignore + return image_with_borders @deprecated( @@ -378,7 +378,7 @@ def grayscale_image(image: ImageType) -> ImageType: ![grayscale-image](https://media.roboflow.com/supervision-docs/supervision-docs-grayscale-image-2.png){ align=center width="1000" } """ # noqa E501 // docs grayscaled = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) - return cv2.cvtColor(grayscaled, cv2.COLOR_GRAY2BGR) # type: ignore + return cv2.cvtColor(grayscaled, cv2.COLOR_GRAY2BGR) def get_image_resolution_wh(image: ImageType) -> tuple[int, int]: