feature/new_video_utils

README.md

@@ -25,7 +25,7 @@
 ## 💻 install
 
 Pip install the supervision package in a
-[**3.10>=Python>=3.7**](https://www.python.org/) environment.
+[**3.11>=Python>=3.7**](https://www.python.org/) environment.
 
 ```bash
 pip install supervision

docs/image.md

@@ -0,0 +1,3 @@
+## ImageSink
+
+:::supervision.image.ImageSink

docs/index.md

@@ -19,7 +19,7 @@ pre-release stage 🚧 Keep your eyes open for potential bugs and be aware that
 ## 💻 How to Install
 
 You can install `supervision` with pip in a 
-[**3.10>=Python>=3.7**](https://www.python.org/) environment.
+[**3.11>=Python>=3.7**](https://www.python.org/) environment.
 
 !!! example "Pip install method (recommended)"
 

mkdocs.yml

@@ -43,6 +43,8 @@ nav:
         - Utils: notebook/utils.md
     - Video:
         - Utils: video.md
+    - Image:
+        - Utils: image.md
   - Changelog: changelog.md
 
 theme:

setup.py

@@ -52,6 +52,7 @@ def get_version():
         'Programming Language :: Python :: 3.8',
         'Programming Language :: Python :: 3.9',
         'Programming Language :: Python :: 3.10',
+        'Programming Language :: Python :: 3.11',
         'Programming Language :: Python :: 3 :: Only',
         'Topic :: Software Development',
         'Topic :: Scientific/Engineering',

supervision/__init__.py

@@ -19,6 +19,7 @@
 from supervision.file import list_files_with_extensions
 from supervision.geometry.core import Point, Position, Rect
 from supervision.geometry.utils import get_polygon_center
+from supervision.image import ImageSink
 from supervision.notebook.utils import plot_image, plot_images_grid
 from supervision.video import (
     VideoInfo,

supervision/detection/core.py

@@ -3,9 +3,14 @@
 from dataclasses import astuple, dataclass
 from typing import Any, Iterator, List, Optional, Tuple, Union
 
-import numpy as np
 import cv2
-from supervision.detection.utils import non_max_suppression, xywh_to_xyxy
+import numpy as np
+
+from supervision.detection.utils import (
+    extract_yolov8_masks,
+    non_max_suppression,
+    xywh_to_xyxy,
+)
 from supervision.geometry.core import Position
 from supervision.internal import deprecated
 
@@ -187,33 +192,11 @@ def from_yolov8(cls, yolov8_results) -> Detections:
             >>> detections = sv.Detections.from_yolov8(result)
             ```
         """
-        mask_maps = None
-        if yolov8_results.masks:
-            orig_shape = yolov8_results.orig_shape
-            inference_shape = tuple(yolov8_results.masks.data.shape[1:])
-            masks = yolov8_results.masks.data.cpu().numpy()
-            # calculate pad and gain
-            pad = (0, 0)
-            gain = 0
-            if inference_shape != orig_shape:
-                gain = min(inference_shape[0] / orig_shape[0], inference_shape[1] / orig_shape[1])  # gain  = old / new
-                pad = (inference_shape[1] - orig_shape[1] * gain) / 2, (inference_shape[0] - orig_shape[0] * gain) / 2
-
-            top, left = int(pad[1]), int(pad[0])  # y, x
-            bottom, right = int(inference_shape[0] - pad[1]), int(inference_shape[1] - pad[0])
-
-            mask_maps = []
-            for i in range(masks.shape[0]):
-                mask = masks[i]
-                mask = mask[top:bottom, left:right]
-                mask = cv2.resize(mask, (orig_shape[1], orig_shape[0]))
-                mask_maps.append(mask)
-            mask_maps = np.asarray(mask_maps)
         return cls(
             xyxy=yolov8_results.boxes.xyxy.cpu().numpy(),
             confidence=yolov8_results.boxes.conf.cpu().numpy(),
             class_id=yolov8_results.boxes.cls.cpu().numpy().astype(int),
-            mask=mask_maps
+            mask=extract_yolov8_masks(yolov8_results),
         )
 
     @classmethod

supervision/detection/utils.py

@@ -258,3 +258,37 @@ def approximate_polygon(
             break
 
     return np.squeeze(approximated_points, axis=1)
+
+
+def extract_yolov8_masks(yolov8_results) -> Optional[np.ndarray]:
+    if not yolov8_results.masks:
+        return None
+
+    orig_shape = yolov8_results.orig_shape
+    inference_shape = tuple(yolov8_results.masks.data.shape[1:])
+
+    gain = 0
+    pad = (0, 0)
+
+    if inference_shape != orig_shape:
+        gain = min(
+            inference_shape[0] / orig_shape[0],
+            inference_shape[1] / orig_shape[1],
+        )
+        pad = (
+            (inference_shape[1] - orig_shape[1] * gain) / 2,
+            (inference_shape[0] - orig_shape[0] * gain) / 2,
+        )
+
+    top, left = int(pad[1]), int(pad[0])
+    bottom, right = int(inference_shape[0] - pad[1]), int(inference_shape[1] - pad[0])
+
+    mask_maps = []
+    masks = yolov8_results.masks.data.cpu().numpy()
+    for i in range(masks.shape[0]):
+        mask = masks[i]
+        mask = mask[top:bottom, left:right]
+        mask = cv2.resize(mask, (orig_shape[1], orig_shape[0]))
+        mask_maps.append(mask)
+
+    return np.asarray(mask_maps, dtype=bool)

supervision/image.py

@@ -0,0 +1,64 @@
+import os
+import shutil
+from typing import Optional
+
+import cv2
+import numpy as np
+
+
+class ImageSink:
+    def __init__(
+        self,
+        target_dir_path: str,
+        overwrite: bool = False,
+        image_name_pattern: str = "image_{:05d}.png",
+    ):
+        """
+        Initialize a context manager for saving images.
+
+        Args:
+            target_dir_path (str): The target directory where images will be saved.
+            overwrite (bool, optional): Whether to overwrite the existing directory. Defaults to False.
+            image_name_pattern (str, optional): The image file name pattern. Defaults to "image_{:05d}.png".
+
+        Example:
+            ```python
+            >>> import supervision as sv
+
+            >>> with sv.ImageSink(target_dir_path='target/directory/path', overwrite=True) as sink:
+            ...     for image in sv.get_video_frames_generator(source_path='source_video.mp4', stride=2):
+            ...         sink.save_image(image=image)
+            ```
+        """
+        self.target_dir_path = target_dir_path
+        self.overwrite = overwrite
+        self.image_name_pattern = image_name_pattern
+        self.image_count = 0
+
+    def __enter__(self):
+        if os.path.exists(self.target_dir_path):
+            if self.overwrite:
+                shutil.rmtree(self.target_dir_path)
+                os.makedirs(self.target_dir_path)
+        else:
+            os.makedirs(self.target_dir_path)
+
+        return self
+
+    def save_image(self, image: np.ndarray, image_name: Optional[str] = None):
+        """
+        Save a given image in the target directory.
+
+        Args:
+            image (np.ndarray): The image to be saved.
+            image_name (str, optional): The name to use for the saved image. If not provided, a name will be generated using the `image_name_pattern`.
+        """
+        if image_name is None:
+            image_name = self.image_name_pattern.format(self.image_count)
+
+        image_path = os.path.join(self.target_dir_path, image_name)
+        cv2.imwrite(image_path, image)
+        self.image_count += 1
+
+    def __exit__(self, exc_type, exc_value, exc_traceback):
+        pass

supervision/video.py

@@ -20,9 +20,9 @@ class VideoInfo:
 
     Examples:
         ```python
-        >>> from supervision import VideoInfo
+        >>> import supervision as sv
 
-        >>> video_info = VideoInfo.from_video_path(video_path='video.mp4')
+        >>> video_info = sv.VideoInfo.from_video_path(video_path='video.mp4')
 
         >>> video_info
         VideoInfo(width=3840, height=2160, fps=25, total_frames=538)
@@ -63,15 +63,15 @@ class VideoSink:
         target_path (str): The path to the output file where the video will be saved.
         video_info (VideoInfo): Information about the video resolution, fps, and total frame count.
 
-    Examples:
+    Example:
         ```python
-        >>> from supervision import VideoInfo, VideoSink
+        >>> import supervision as sv
 
-        >>> video_info = VideoInfo.from_video_path(video_path='source_video.mp4')
+        >>> video_info = sv.VideoInfo.from_video_path(video_path='source_video.mp4')
 
-        >>> with VideoSink(target_path='target_video.mp4', video_info=video_info) as s:
-        ...     frame = ...
-        ...     s.write_frame(frame=frame)
+        >>> with sv.VideoSink(target_path='target_video.mp4', video_info=video_info) as sink:
+        ...     for frame in get_video_frames_generator(source_path='source_video.mp4', stride=2):
+        ...         sink.write_frame(frame=frame)
         ```
     """
 
@@ -93,35 +93,43 @@ def __enter__(self):
     def write_frame(self, frame: np.ndarray):
         self.__writer.write(frame)
 
-    def __exit__(self, exc_type, exc_val, exc_tb):
+    def __exit__(self, exc_type, exc_value, exc_traceback):
         self.__writer.release()
 
 
-def get_video_frames_generator(source_path: str) -> Generator[np.ndarray, None, None]:
+def get_video_frames_generator(
+    source_path: str, stride: int = 1
+) -> Generator[np.ndarray, None, None]:
     """
     Get a generator that yields the frames of the video.
 
     Args:
         source_path (str): The path of the video file.
+        stride (int): The number of frames to skip before returning the next one.
 
     Returns:
         (Generator[np.ndarray, None, None]): A generator that yields the frames of the video.
 
     Examples:
         ```python
-        >>> from supervision import get_video_frames_generator
+        >>> import supervision as sv
 
-        >>> for frame in get_video_frames_generator(source_path='source_video.mp4'):
+        >>> for frame in sv.get_video_frames_generator(source_path='source_video.mp4', stride=2):
         ...     ...
         ```
     """
     video = cv2.VideoCapture(source_path)
     if not video.isOpened():
         raise Exception(f"Could not open video at {source_path}")
+
+    frame_count = 0
     success, frame = video.read()
     while success:
-        yield frame
+        if frame_count % stride == 0:
+            yield frame
         success, frame = video.read()
+        frame_count += 1
+
     video.release()