diff --git a/supervision/draw/utils.py b/supervision/draw/utils.py index ae7d8c69..a49605ad 100644 --- a/supervision/draw/utils.py +++ b/supervision/draw/utils.py @@ -1,8 +1,8 @@ import cv2 import numpy as np -from supervision.geometry.dataclasses import Point, Rect from supervision.draw.color import Color +from supervision.geometry.dataclasses import Point, Rect def draw_line( @@ -95,7 +95,9 @@ def draw_filled_rectangle(scene: np.ndarray, rect: Rect, color: Color) -> np.nda return scene -def draw_polygon(scene: np.ndarray, polygon: np.ndarray, color: Color, thickness: int = 2) -> np.ndarray: +def draw_polygon( + scene: np.ndarray, polygon: np.ndarray, color: Color, thickness: int = 2 +) -> np.ndarray: """Draw a polygon on a scene. Attributes: @@ -107,5 +109,7 @@ def draw_polygon(scene: np.ndarray, polygon: np.ndarray, color: Color, thickness Returns: np.ndarray: The scene with the polygon drawn on it. """ - cv2.polylines(scene, [polygon], isClosed=True, color=color.as_bgr(), thickness=thickness) + cv2.polylines( + scene, [polygon], isClosed=True, color=color.as_bgr(), thickness=thickness + ) return scene diff --git a/supervision/geometry/dataclasses.py b/supervision/geometry/dataclasses.py index 017ebcc8..fc352f55 100644 --- a/supervision/geometry/dataclasses.py +++ b/supervision/geometry/dataclasses.py @@ -1,13 +1,13 @@ from __future__ import annotations from dataclasses import dataclass -from typing import Tuple from enum import Enum +from typing import Tuple class Position(Enum): - CENTER = 'CENTER' - BOTTOM_CENTER = 'BOTTOM_CENTER' + CENTER = "CENTER" + BOTTOM_CENTER = "BOTTOM_CENTER" @classmethod def list(cls): diff --git a/supervision/tools/detections.py b/supervision/tools/detections.py index 8f97d930..79cb2cb8 100644 --- a/supervision/tools/detections.py +++ b/supervision/tools/detections.py @@ -128,15 +128,16 @@ class Detections: np.ndarray: An array of shape (n, 2) containing the bounding box anchor coordinates in format [x, y]. """ if anchor == Position.CENTER: - return np.array([ - (self.xyxy[:, 0] + self.xyxy[:, 2]) / 2, - (self.xyxy[:, 1] + self.xyxy[:, 3]) / 2 - ]).transpose() + return np.array( + [ + (self.xyxy[:, 0] + self.xyxy[:, 2]) / 2, + (self.xyxy[:, 1] + self.xyxy[:, 3]) / 2, + ] + ).transpose() elif anchor == Position.BOTTOM_CENTER: - return np.array([ - (self.xyxy[:, 0] + self.xyxy[:, 2]) / 2, - self.xyxy[:, 3] - ]).transpose() + return np.array( + [(self.xyxy[:, 0] + self.xyxy[:, 2]) / 2, self.xyxy[:, 3]] + ).transpose() raise ValueError(f"{anchor} is not supported.") diff --git a/supervision/video.py b/supervision/video.py index 92ad0a6e..228b5a2d 100644 --- a/supervision/video.py +++ b/supervision/video.py @@ -1,11 +1,13 @@ from __future__ import annotations +from dataclasses import dataclass from typing import Callable, Generator, Optional, Tuple import cv2 import numpy as np +@dataclass class VideoInfo: """ A class to store video information, including width, height, fps and total number of frames. @@ -24,16 +26,16 @@ class VideoInfo: >>> video_info VideoInfo(width=3840, height=2160, fps=25, total_frames=538) + + >>> video_info.resolution_wh + (3840, 2160) ``` """ - def __init__( - self, width: int, height: int, fps: int, total_frames: Optional[int] = None - ): - self.width = width - self.height = height - self.fps = fps - self.total_frames = total_frames + width: int + height: int + fps: int + total_frames: Optional[int] = None @classmethod def from_video_path(cls, video_path: str) -> VideoInfo: @@ -49,7 +51,7 @@ class VideoInfo: return VideoInfo(width, height, fps, total_frames) @property - def resolution(self) -> Tuple[int, int]: + def resolution_wh(self) -> Tuple[int, int]: return self.width, self.height