From 98f60c24d3ec0baa06150daeef70f3bc3b4ccd05 Mon Sep 17 00:00:00 2001
From: SkalskiP <piotr.skalski92@gmail.com>
Date: Tue, 6 Jun 2023 17:35:40 +0200
Subject: [PATCH] =?UTF-8?q?=F0=9F=98=B7=20Add=20`stride`=20to=20`get=5Fvid?=
 =?UTF-8?q?eo=5Fframes=5Fgenerator`=20function.=20Initial=20`from=5Fyolov8?=
 =?UTF-8?q?`=20with=20`masks`=20refactor.?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 supervision/detection/core.py | 18 +++++++++++++-----
 supervision/video.py          | 17 +++++++++++++----
 2 files changed, 26 insertions(+), 9 deletions(-)

diff --git a/supervision/detection/core.py b/supervision/detection/core.py
index 60809779..b6845be5 100644
--- a/supervision/detection/core.py
+++ b/supervision/detection/core.py
@@ -3,8 +3,9 @@ from __future__ import annotations
 from dataclasses import astuple, dataclass
 from typing import Any, Iterator, List, Optional, Tuple, Union
 
-import numpy as np
 import cv2
+import numpy as np
+
 from supervision.detection.utils import non_max_suppression, xywh_to_xyxy
 from supervision.geometry.core import Position
 from supervision.internal import deprecated
@@ -196,11 +197,18 @@ class Detections:
             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
+                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])
+            bottom, right = int(inference_shape[0] - pad[1]), int(
+                inference_shape[1] - pad[0]
+            )
 
             mask_maps = []
             for i in range(masks.shape[0]):
@@ -213,7 +221,7 @@ class Detections:
             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=mask_maps,
         )
 
     @classmethod
diff --git a/supervision/video.py b/supervision/video.py
index 5dd807eb..0e147536 100644
--- a/supervision/video.py
+++ b/supervision/video.py
@@ -97,31 +97,40 @@ class VideoSink:
         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()