Update VideoInfo API 🎬

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

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):

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.")
 

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