From 301f1df3868f018e2a9734c6cca9c2b2a14a81f8 Mon Sep 17 00:00:00 2001 From: SkalskiP Date: Mon, 6 Feb 2023 13:55:57 +0100 Subject: [PATCH 01/23] =?UTF-8?q?=F0=9F=8E=89=20initial=20commit=200.2.0?= =?UTF-8?q?=20API=20commit?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- supervision/draw/utils.py | 18 +++++++++++++++++- supervision/geometry/dataclasses.py | 10 ++++++++++ supervision/notebook/utils.py | 2 +- supervision/tools/detections.py | 24 ++++++++++++++++++++++++ supervision/tools/polygon_zone.py | 0 5 files changed, 52 insertions(+), 2 deletions(-) create mode 100644 supervision/tools/polygon_zone.py diff --git a/supervision/draw/utils.py b/supervision/draw/utils.py index ff6c050c7..ae7d8c690 100644 --- a/supervision/draw/utils.py +++ b/supervision/draw/utils.py @@ -1,7 +1,7 @@ import cv2 import numpy as np -from supervision.commons.dataclasses import Point, Rect +from supervision.geometry.dataclasses import Point, Rect from supervision.draw.color import Color @@ -93,3 +93,19 @@ def draw_filled_rectangle(scene: np.ndarray, rect: Rect, color: Color) -> np.nda -1, ) return scene + + +def draw_polygon(scene: np.ndarray, polygon: np.ndarray, color: Color, thickness: int = 2) -> np.ndarray: + """Draw a polygon on a scene. + + Attributes: + scene (np.ndarray): The scene to draw the polygon on. + polygon (np.ndarray): The polygon to be drawn, given as a list of vertices. + color (Color): The color of the polygon. + thickness (int, optional): The thickness of the polygon lines, by default 2. + + Returns: + np.ndarray: The scene with the polygon drawn on it. + """ + 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 b647ec096..017ebcc80 100644 --- a/supervision/geometry/dataclasses.py +++ b/supervision/geometry/dataclasses.py @@ -2,6 +2,16 @@ from dataclasses import dataclass from typing import Tuple +from enum import Enum + + +class Position(Enum): + CENTER = 'CENTER' + BOTTOM_CENTER = 'BOTTOM_CENTER' + + @classmethod + def list(cls): + return list(map(lambda c: c.value, cls)) @dataclass diff --git a/supervision/notebook/utils.py b/supervision/notebook/utils.py index c50c98428..4e5d52297 100644 --- a/supervision/notebook/utils.py +++ b/supervision/notebook/utils.py @@ -18,7 +18,7 @@ def show_frame_in_notebook( Examples: ```python - >>> from supervision.notebook import show_frame_in_notebook + >>> from supervision.notebook.utils import show_frame_in_notebook ``` """ diff --git a/supervision/tools/detections.py b/supervision/tools/detections.py index 4faaf6cf6..8f97d9303 100644 --- a/supervision/tools/detections.py +++ b/supervision/tools/detections.py @@ -4,6 +4,7 @@ import numpy as np from supervision.draw.color import Color, ColorPalette +from supervision.geometry.dataclasses import Position class Detections: @@ -116,6 +117,29 @@ def filter(self, mask: np.ndarray, inplace: bool = False) -> Optional[np.ndarray else None, ) + def get_anchor_coordinates(self, anchor: Position) -> np.ndarray: + """ + Returns the bounding box coordinates for a specific anchor. + + Attributes: + anchor (Position): Position of bounding box anchor for which to return the coordinates. + + Returns: + 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() + elif anchor == Position.BOTTOM_CENTER: + return np.array([ + (self.xyxy[:, 0] + self.xyxy[:, 2]) / 2, + self.xyxy[:, 3] + ]).transpose() + + raise ValueError(f"{anchor} is not supported.") + class BoxAnnotator: def __init__( diff --git a/supervision/tools/polygon_zone.py b/supervision/tools/polygon_zone.py new file mode 100644 index 000000000..e69de29bb From 0e3851506cd21ea6eaf7c896ef4e0591072a2af3 Mon Sep 17 00:00:00 2001 From: SkalskiP Date: Mon, 6 Feb 2023 14:20:56 +0100 Subject: [PATCH 02/23] =?UTF-8?q?Update=20VideoInfo=20API=20=F0=9F=8E=AC?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- supervision/draw/utils.py | 10 +++++++--- supervision/geometry/dataclasses.py | 6 +++--- supervision/tools/detections.py | 17 +++++++++-------- supervision/video.py | 18 ++++++++++-------- 4 files changed, 29 insertions(+), 22 deletions(-) diff --git a/supervision/draw/utils.py b/supervision/draw/utils.py index ae7d8c690..a49605adb 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 017ebcc80..fc352f558 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 8f97d9303..79cb2cb88 100644 --- a/supervision/tools/detections.py +++ b/supervision/tools/detections.py @@ -128,15 +128,16 @@ def get_anchor_coordinates(self, anchor: Position) -> np.ndarray: 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 92ad0a6e2..228b5a2d1 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 @@ def from_video_path(cls, video_path: str) -> 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 From d5d7187fe4ce66cfc56b123599a97f8387936d79 Mon Sep 17 00:00:00 2001 From: SkalskiP Date: Mon, 6 Feb 2023 14:28:16 +0100 Subject: [PATCH 03/23] Detection core API reformat --- supervision/draw/utils.py | 2 +- supervision/geometry/{dataclasses.py => core.py} | 0 supervision/tools/{detections.py => core.py} | 2 +- supervision/tools/line_counter.py | 4 ++-- test/geometry/test_dataclasses.py | 2 +- 5 files changed, 5 insertions(+), 5 deletions(-) rename supervision/geometry/{dataclasses.py => core.py} (100%) rename supervision/tools/{detections.py => core.py} (99%) diff --git a/supervision/draw/utils.py b/supervision/draw/utils.py index a49605adb..3cfe3288e 100644 --- a/supervision/draw/utils.py +++ b/supervision/draw/utils.py @@ -2,7 +2,7 @@ import numpy as np from supervision.draw.color import Color -from supervision.geometry.dataclasses import Point, Rect +from supervision.geometry.core import Point, Rect def draw_line( diff --git a/supervision/geometry/dataclasses.py b/supervision/geometry/core.py similarity index 100% rename from supervision/geometry/dataclasses.py rename to supervision/geometry/core.py diff --git a/supervision/tools/detections.py b/supervision/tools/core.py similarity index 99% rename from supervision/tools/detections.py rename to supervision/tools/core.py index 79cb2cb88..edd506b6f 100644 --- a/supervision/tools/detections.py +++ b/supervision/tools/core.py @@ -4,7 +4,7 @@ import numpy as np from supervision.draw.color import Color, ColorPalette -from supervision.geometry.dataclasses import Position +from supervision.geometry.core import Position class Detections: diff --git a/supervision/tools/line_counter.py b/supervision/tools/line_counter.py index f791c9a6f..161954490 100644 --- a/supervision/tools/line_counter.py +++ b/supervision/tools/line_counter.py @@ -4,8 +4,8 @@ import numpy as np from supervision.draw.color import Color -from supervision.geometry.dataclasses import Point, Rect, Vector -from supervision.tools.detections import Detections +from supervision.geometry.core import Point, Rect, Vector +from supervision.tools.core import Detections class LineCounter: diff --git a/test/geometry/test_dataclasses.py b/test/geometry/test_dataclasses.py index 348218846..711407123 100644 --- a/test/geometry/test_dataclasses.py +++ b/test/geometry/test_dataclasses.py @@ -1,6 +1,6 @@ import pytest -from supervision.geometry.dataclasses import Vector, Point +from supervision.geometry.core import Vector, Point @pytest.mark.parametrize( From 307ec8a042b93387439673409481c6613d51cc51 Mon Sep 17 00:00:00 2001 From: SkalskiP Date: Mon, 6 Feb 2023 14:49:09 +0100 Subject: [PATCH 04/23] =?UTF-8?q?=F0=9F=93=84=20Documentation=20updates=20?= =?UTF-8?q?before=200.2.0=20release?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- docs/detection.md | 7 +++ docs/draw.md | 6 +-- docs/notebook.md | 2 +- docs/tools.md | 9 ---- mkdocs.yml | 13 ++--- supervision/{tools => detection}/__init__.py | 0 supervision/{tools => detection}/core.py | 51 +++++++++---------- .../{tools => detection}/line_counter.py | 2 +- .../{tools => detection}/polygon_zone.py | 0 supervision/draw/utils.py | 3 +- supervision/notebook/utils.py | 2 + 11 files changed, 45 insertions(+), 50 deletions(-) create mode 100644 docs/detection.md delete mode 100644 docs/tools.md rename supervision/{tools => detection}/__init__.py (100%) rename supervision/{tools => detection}/core.py (84%) rename supervision/{tools => detection}/line_counter.py (99%) rename supervision/{tools => detection}/polygon_zone.py (100%) diff --git a/docs/detection.md b/docs/detection.md new file mode 100644 index 000000000..cb4119fe5 --- /dev/null +++ b/docs/detection.md @@ -0,0 +1,7 @@ +## Detections + +:::supervision.detection.core.Detections + +## LineCounter + +:::supervision.detection.line_counter.LineCounter \ No newline at end of file diff --git a/docs/draw.md b/docs/draw.md index 8ba499b16..e94509e30 100644 --- a/docs/draw.md +++ b/docs/draw.md @@ -1,9 +1,7 @@ -Utilities for drawing on images. - -## Draw Line +## draw_line :::supervision.draw.utils.draw_line -## Draw Rectangle +## draw_rectangle :::supervision.draw.utils.draw_rectangle \ No newline at end of file diff --git a/docs/notebook.md b/docs/notebook.md index ce650120a..b910496d9 100644 --- a/docs/notebook.md +++ b/docs/notebook.md @@ -1,3 +1,3 @@ -Utilities to help you build computer vision projects in notebook environments. +## show_frame_in_notebook :::supervision.notebook.utils.show_frame_in_notebook \ No newline at end of file diff --git a/docs/tools.md b/docs/tools.md deleted file mode 100644 index 9274cb0ef..000000000 --- a/docs/tools.md +++ /dev/null @@ -1,9 +0,0 @@ -Useful utilities for common computer vision tasks. - -## Helper for Processing Model Detections - -:::supervision.tools.detections.Detections - -## Count Objects That Pass a Line - -:::supervision.tools.line_counter.LineCounter \ No newline at end of file diff --git a/mkdocs.yml b/mkdocs.yml index 3212b2aaf..408e375f6 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -22,12 +22,13 @@ extra: property: G-P7ZG0Y19G5 nav: - - Home 🏠: index.md - - Video 📷: video.md - - Notebook Helpers 📓: notebook.md - - Draw 🎨: draw.md - - Geometry 📐: geometry.md - - Tools 🛠: tools.md + - Home: index.md + - API reference: + - Video: video.md + - Detection: detection.md + - Draw: draw.md + - Geometry: geometry.md + - Notebook: notebook.md theme: name: 'material' diff --git a/supervision/tools/__init__.py b/supervision/detection/__init__.py similarity index 100% rename from supervision/tools/__init__.py rename to supervision/detection/__init__.py diff --git a/supervision/tools/core.py b/supervision/detection/core.py similarity index 84% rename from supervision/tools/core.py rename to supervision/detection/core.py index edd506b6f..c4d5362de 100644 --- a/supervision/tools/core.py +++ b/supervision/detection/core.py @@ -1,3 +1,4 @@ +from dataclasses import dataclass from typing import List, Optional, Union import cv2 @@ -7,28 +8,24 @@ from supervision.geometry.core import Position +@dataclass class Detections: - def __init__( - self, - xyxy: np.ndarray, - confidence: np.ndarray, - class_id: np.ndarray, - tracker_id: Optional[np.ndarray] = None, - ): - """ - Data class containing information about the detections in a video frame. - - Attributes: - xyxy (np.ndarray): An array of shape (n, 4) containing the bounding boxes coordinates in format [x1, y1, x2, y2] - confidence (np.ndarray): An array of shape (n,) containing the confidence scores of the detections. - class_id (np.ndarray): An array of shape (n,) containing the class ids of the detections. - tracker_id (Optional[np.ndarray]): An array of shape (n,) containing the tracker ids of the detections. - """ - self.xyxy: np.ndarray = xyxy - self.confidence: np.ndarray = confidence - self.class_id: np.ndarray = class_id - self.tracker_id: Optional[np.ndarray] = tracker_id - + """ + Data class containing information about the detections in a video frame. + + Attributes: + xyxy (np.ndarray): An array of shape `(n, 4)` containing the bounding boxes coordinates in format `[x1, y1, x2, y2]` + confidence (np.ndarray): An array of shape `(n,)` containing the confidence scores of the detections. + class_id (np.ndarray): An array of shape `(n,)` containing the class ids of the detections. + tracker_id (Optional[np.ndarray]): An array of shape `(n,)` containing the tracker ids of the detections. + """ + + xyxy: np.ndarray + confidence: np.ndarray + class_id: np.ndarray + tracker_id: Optional[np.ndarray] = None + + def __post_init__(self): n = len(self.xyxy) validators = [ (isinstance(self.xyxy, np.ndarray) and self.xyxy.shape == (n, 4)), @@ -56,7 +53,7 @@ def __len__(self): def __iter__(self): """ - Iterates over the Detections object and yield a tuple of (xyxy, confidence, class_id, tracker_id) for each detection. + Iterates over the Detections object and yield a tuple of `(xyxy, confidence, class_id, tracker_id)` for each detection. """ for i in range(len(self.xyxy)): yield ( @@ -78,7 +75,7 @@ def from_yolov5(cls, yolov5_output: np.ndarray): Example: ```python - >>> from supervision.tools.detections import Detections + >>> from supervision.detection.core import Detections >>> detections = Detections.from_yolov5(yolov5_output) ``` @@ -93,11 +90,11 @@ def filter(self, mask: np.ndarray, inplace: bool = False) -> Optional[np.ndarray Filter the detections by applying a mask. Attributes: - mask (np.ndarray): A mask of shape (n,) containing a boolean value for each detection indicating if it should be included in the filtered detections + mask (np.ndarray): A mask of shape `(n,)` containing a boolean value for each detection indicating if it should be included in the filtered detections inplace (bool): If True, the original data will be modified and self will be returned. Returns: - Optional[np.ndarray]: A new instance of Detections with the filtered detections, if inplace is set to False. None otherwise. + Optional[np.ndarray]: A new instance of Detections with the filtered detections, if inplace is set to `False`. `None` otherwise. """ if inplace: self.xyxy = self.xyxy[mask] @@ -121,11 +118,11 @@ def get_anchor_coordinates(self, anchor: Position) -> np.ndarray: """ Returns the bounding box coordinates for a specific anchor. - Attributes: + Properties: anchor (Position): Position of bounding box anchor for which to return the coordinates. Returns: - np.ndarray: An array of shape (n, 2) containing the bounding box anchor coordinates in format [x, y]. + 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( diff --git a/supervision/tools/line_counter.py b/supervision/detection/line_counter.py similarity index 99% rename from supervision/tools/line_counter.py rename to supervision/detection/line_counter.py index 161954490..030508cae 100644 --- a/supervision/tools/line_counter.py +++ b/supervision/detection/line_counter.py @@ -3,9 +3,9 @@ import cv2 import numpy as np +from supervision.detection.core import Detections from supervision.draw.color import Color from supervision.geometry.core import Point, Rect, Vector -from supervision.tools.core import Detections class LineCounter: diff --git a/supervision/tools/polygon_zone.py b/supervision/detection/polygon_zone.py similarity index 100% rename from supervision/tools/polygon_zone.py rename to supervision/detection/polygon_zone.py diff --git a/supervision/draw/utils.py b/supervision/draw/utils.py index 3cfe3288e..d4084c542 100644 --- a/supervision/draw/utils.py +++ b/supervision/draw/utils.py @@ -11,8 +11,7 @@ def draw_line( """ Draws a line on a given scene. - Attributes: - + Parameters: scene (np.ndarray): The scene on which the line will be drawn start (Point): The starting point of the line end (Point): The end point of the line diff --git a/supervision/notebook/utils.py b/supervision/notebook/utils.py index 4e5d52297..a750daa17 100644 --- a/supervision/notebook/utils.py +++ b/supervision/notebook/utils.py @@ -20,6 +20,8 @@ def show_frame_in_notebook( ```python >>> from supervision.notebook.utils import show_frame_in_notebook + %matplotlib inline + show_frame_in_notebook(frame, (16, 16)) ``` """ if frame.ndim == 2: From 0ac184b01d90e0cc8f5b6d1c564345fa09dfb6f5 Mon Sep 17 00:00:00 2001 From: SkalskiP Date: Mon, 6 Feb 2023 15:23:20 +0100 Subject: [PATCH 05/23] =?UTF-8?q?=F0=9F=91=8A=20bump=20Supervision=20versi?= =?UTF-8?q?on=20to=200.2.0.dev0?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- supervision/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/supervision/__init__.py b/supervision/__init__.py index 3dc1f76bc..b9d465bd8 100644 --- a/supervision/__init__.py +++ b/supervision/__init__.py @@ -1 +1 @@ -__version__ = "0.1.0" +__version__ = "0.2.0.dev0" From b7c1791236d5a9cffd23e2d1af5091fa60774f50 Mon Sep 17 00:00:00 2001 From: SkalskiP Date: Mon, 6 Feb 2023 16:14:24 +0100 Subject: [PATCH 06/23] =?UTF-8?q?=F0=9F=95=B5=EF=B8=8F=20update=20Detectio?= =?UTF-8?q?n=20API?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- docs/detection.md | 7 ------- docs/detection_core.md | 3 +++ docs/detection_utils.md | 3 +++ mkdocs.yml | 4 +++- supervision/detection/core.py | 18 +++++++++++------- supervision/detection/utils.py | 20 ++++++++++++++++++++ 6 files changed, 40 insertions(+), 15 deletions(-) delete mode 100644 docs/detection.md create mode 100644 docs/detection_core.md create mode 100644 docs/detection_utils.md create mode 100644 supervision/detection/utils.py diff --git a/docs/detection.md b/docs/detection.md deleted file mode 100644 index cb4119fe5..000000000 --- a/docs/detection.md +++ /dev/null @@ -1,7 +0,0 @@ -## Detections - -:::supervision.detection.core.Detections - -## LineCounter - -:::supervision.detection.line_counter.LineCounter \ No newline at end of file diff --git a/docs/detection_core.md b/docs/detection_core.md new file mode 100644 index 000000000..fc3814fa9 --- /dev/null +++ b/docs/detection_core.md @@ -0,0 +1,3 @@ +## Detections + +:::supervision.detection.core.Detections diff --git a/docs/detection_utils.md b/docs/detection_utils.md new file mode 100644 index 000000000..500f8b95b --- /dev/null +++ b/docs/detection_utils.md @@ -0,0 +1,3 @@ +## generate_2d_mask + +:::supervision.detection.utils.generate_2d_mask \ No newline at end of file diff --git a/mkdocs.yml b/mkdocs.yml index 408e375f6..b788d0ce4 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -25,7 +25,9 @@ nav: - Home: index.md - API reference: - Video: video.md - - Detection: detection.md + - Detection: + - Core: detection_core.md + - Utils: detection_utils.md - Draw: draw.md - Geometry: geometry.md - Notebook: notebook.md diff --git a/supervision/detection/core.py b/supervision/detection/core.py index c4d5362de..e83d8fd62 100644 --- a/supervision/detection/core.py +++ b/supervision/detection/core.py @@ -64,25 +64,29 @@ def __iter__(self): ) @classmethod - def from_yolov5(cls, yolov5_output: np.ndarray): + def from_yolov5(cls, yolov5_detections: "yolov5.models.common.Detections"): """ - Creates a Detections instance from a YOLOv5 output tensor + Creates a Detections instance from a YOLOv5 output Detections Attributes: - yolov5_output (np.ndarray): The output tensor from YOLOv5 + yolov5_detections (yolov5.models.common.Detections): The output Detections instance from YOLOv5 Returns: Example: ```python + >>> import torch >>> from supervision.detection.core import Detections - >>> detections = Detections.from_yolov5(yolov5_output) + >>> model = torch.hub.load('ultralytics/yolov5', 'yolov5s') + >>> results = model(frame) + >>> detections = Detections.from_yolov5(results) ``` """ - xyxy = yolov5_output[:, :4] - confidence = yolov5_output[:, 4] - class_id = yolov5_output[:, 5].astype(int) + yolov5_detections_predictions = yolov5_detections.pred[0].cpu().cpu().numpy() + xyxy = yolov5_detections_predictions[:, :4] + confidence = yolov5_detections_predictions[:, 4] + class_id = yolov5_detections_predictions[:, 5].astype(int) return cls(xyxy, confidence, class_id) def filter(self, mask: np.ndarray, inplace: bool = False) -> Optional[np.ndarray]: diff --git a/supervision/detection/utils.py b/supervision/detection/utils.py new file mode 100644 index 000000000..a3a33490f --- /dev/null +++ b/supervision/detection/utils.py @@ -0,0 +1,20 @@ +from typing import Tuple + +import cv2 +import numpy as np + + +def generate_2d_mask(polygon: np.ndarray, resolution_wh: Tuple[int, int]) -> np.ndarray: + """Generate a 2D mask from a polygon. + + Properties: + polygon (np.ndarray): The polygon for which the mask should be generated, given as a list of vertices. + resolution_wh (Tuple[int, int]): The width and height of the desired resolution. + + Returns: + np.ndarray: The generated 2D mask, where the polygon is marked with `1`'s and the rest is filled with `0`'s. + """ + width, height = resolution_wh + mask = np.zeros((height, width), dtype=np.uint8) + cv2.fillPoly(mask, [polygon], color=1) + return mask From 66f81ceeb9a333779fd508c62650a0d9cd25c7be Mon Sep 17 00:00:00 2001 From: SkalskiP Date: Mon, 6 Feb 2023 22:26:44 +0100 Subject: [PATCH 07/23] - ColorPalette.default() added - Detections.from_yolov5() updated - Detection.__getitem__ added - get_polygon_center added --- supervision/__init__.py | 2 ++ supervision/detection/core.py | 18 +++++++++++++++--- supervision/draw/color.py | 12 ++++++------ supervision/geometry/utils.py | 28 ++++++++++++++++++++++++++++ test/detection/__init__.py | 0 5 files changed, 51 insertions(+), 9 deletions(-) create mode 100644 supervision/geometry/utils.py create mode 100644 test/detection/__init__.py diff --git a/supervision/__init__.py b/supervision/__init__.py index b9d465bd8..9a0f17def 100644 --- a/supervision/__init__.py +++ b/supervision/__init__.py @@ -1 +1,3 @@ __version__ = "0.2.0.dev0" + +from supervision.detection.core import Detections diff --git a/supervision/detection/core.py b/supervision/detection/core.py index e83d8fd62..2fef419b5 100644 --- a/supervision/detection/core.py +++ b/supervision/detection/core.py @@ -1,3 +1,5 @@ +from __future__ import annotations + from dataclasses import dataclass from typing import List, Optional, Union @@ -76,7 +78,7 @@ def from_yolov5(cls, yolov5_detections: "yolov5.models.common.Detections"): Example: ```python >>> import torch - >>> from supervision.detection.core import Detections + >>> from supervision import Detections >>> model = torch.hub.load('ultralytics/yolov5', 'yolov5s') >>> results = model(frame) @@ -89,7 +91,7 @@ def from_yolov5(cls, yolov5_detections: "yolov5.models.common.Detections"): class_id = yolov5_detections_predictions[:, 5].astype(int) return cls(xyxy, confidence, class_id) - def filter(self, mask: np.ndarray, inplace: bool = False) -> Optional[np.ndarray]: + def filter(self, mask: np.ndarray, inplace: bool = False) -> Optional[Detections]: """ Filter the detections by applying a mask. @@ -142,11 +144,21 @@ def get_anchor_coordinates(self, anchor: Position) -> np.ndarray: raise ValueError(f"{anchor} is not supported.") + def __getitem__(self, index: np.ndarray) -> Detections: + if isinstance(index, np.ndarray) and index.dtype == np.bool: + return Detections( + xyxy=self.xyxy[index], + confidence=self.confidence[index], + class_id=self.class_id[index], + tracker_id=self.tracker_id[index] if self.tracker_id is not None else None, + ) + raise TypeError(f"Detections.__getitem__ not supported for index of type {type(index)}.") + class BoxAnnotator: def __init__( self, - color: Union[Color, ColorPalette], + color: Union[Color, ColorPalette] = ColorPalette.default(), thickness: int = 2, text_color: Color = Color.black(), text_scale: float = 0.5, diff --git a/supervision/draw/color.py b/supervision/draw/color.py index 842dd2be1..b597831ea 100644 --- a/supervision/draw/color.py +++ b/supervision/draw/color.py @@ -1,6 +1,6 @@ from __future__ import annotations -from dataclasses import dataclass, field +from dataclasses import dataclass from typing import List, Tuple DEFAULT_COLOR_PALETTE = [ @@ -94,11 +94,11 @@ def blue(cls) -> Color: @dataclass class ColorPalette: - colors: List[Color] = field( - default_factory=lambda: [ - Color.from_hex(color_hex) for color_hex in DEFAULT_COLOR_PALETTE - ] - ) + colors: List[Color] + + @classmethod + def default(cls) -> ColorPalette: + return ColorPalette.from_hex(color_hex_list=DEFAULT_COLOR_PALETTE) @classmethod def from_hex(cls, color_hex_list: List[str]): diff --git a/supervision/geometry/utils.py b/supervision/geometry/utils.py new file mode 100644 index 000000000..400e4d755 --- /dev/null +++ b/supervision/geometry/utils.py @@ -0,0 +1,28 @@ +import numpy as np + +from supervision.geometry.core import Point + + +def get_polygon_center(polygon: np.ndarray) -> Point: + """ + Calculate the center of a polygon. + + This function takes in a polygon as a 2-dimensional numpy ndarray and returns the center of the polygon as a Point object. The center is calculated as the mean of the polygon's vertices along each axis, and is rounded down to the nearest integer. + + Parameters: + polygon (np.ndarray): A 2-dimensional numpy ndarray representing the vertices of the polygon. + + Returns: + Point: The center of the polygon, represented as a Point object with x and y attributes. + + Examples: + ```python + >>> from supervision.geometry.utils import get_polygon_center + + >>> vertices = np.array([[0, 0], [0, 1], [1, 1], [1, 0]]) + >>> get_center(vertices) + Point(x=0.5, y=0.5) + ``` + """ + center = np.mean(polygon, axis=0).astype(int) + return Point(x=center[0], y=center[1]) diff --git a/test/detection/__init__.py b/test/detection/__init__.py new file mode 100644 index 000000000..e69de29bb From 807785a8206b56d15438ea2c462e6eac85a8897b Mon Sep 17 00:00:00 2001 From: SkalskiP Date: Mon, 6 Feb 2023 22:35:00 +0100 Subject: [PATCH 08/23] - draw_text added --- supervision/draw/utils.py | 79 +++++++++++++++++++++++++++++++++------ 1 file changed, 68 insertions(+), 11 deletions(-) diff --git a/supervision/draw/utils.py b/supervision/draw/utils.py index d4084c542..ca9366de9 100644 --- a/supervision/draw/utils.py +++ b/supervision/draw/utils.py @@ -4,6 +4,8 @@ from supervision.draw.color import Color from supervision.geometry.core import Point, Rect +from typing import Optional + def draw_line( scene: np.ndarray, start: Point, end: Point, color: Color, thickness: int = 2 @@ -45,11 +47,6 @@ def draw_rectangle( Returns: np.ndarray: The scene with the rectangle drawn on it - - Example: - ```python - >>> # TODO: Add example - ``` """ cv2.rectangle( scene, @@ -77,12 +74,6 @@ def draw_filled_rectangle(scene: np.ndarray, rect: Rect, color: Color) -> np.nda Returns: np.ndarray: The scene with the rectangle drawn on it - - Example: - ```python - >>> # TODO: Add example - ``` - """ cv2.rectangle( scene, @@ -112,3 +103,69 @@ def draw_polygon( scene, [polygon], isClosed=True, color=color.as_bgr(), thickness=thickness ) return scene + + +def draw_text( + scene: np.ndarray, + text: str, + text_anchor: Point, + text_color: Color = Color.black(), + text_scale: float = 0.5, + text_thickness: int = 1, + text_padding: int = 10, + text_font: int = cv2.FONT_HERSHEY_SIMPLEX, + background_color: Optional[Color] = None, +) -> np.ndarray: + """ + Draw text on a scene. + + This function takes in a 2-dimensional numpy ndarray representing an image or scene, and draws text on it using OpenCV's putText function. The text is anchored at a specified Point, and its appearance can be customized using arguments such as color, scale, and font. An optional background color and padding can be specified to draw a rectangle behind the text. + + Parameters: + scene (np.ndarray): A 2-dimensional numpy ndarray representing an image or scene. + text (str): The text to be drawn. + text_anchor (Point): The anchor point for the text, represented as a Point object with x and y attributes. + text_color (Color, optional): The color of the text. Defaults to black. + text_scale (float, optional): The scale of the text. Defaults to 0.5. + text_thickness (int, optional): The thickness of the text. Defaults to 1. + text_padding (int, optional): The amount of padding to add around the text when drawing a rectangle in the background. Defaults to 10. + text_font (int, optional): The font to use for the text. Defaults to cv2.FONT_HERSHEY_SIMPLEX. + background_color (Color, optional): The color of the background rectangle, if one is to be drawn. Defaults to None. + + Returns: + np.ndarray: The input scene with the text drawn on it. + + Examples: + ```python + >>> 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) + ``` + """ + text_width, text_height = cv2.getTextSize( + text=text, + fontFace=text_font, + fontScale=text_scale, + thickness=text_thickness, + )[0] + text_rect = Rect( + x=text_anchor.x - text_width // 2, + y=text_anchor.y - text_height // 2, + width=text_width, + height=text_height + ).pad(text_padding) + + if background_color is not None: + scene = draw_filled_rectangle(scene=scene, rect=text_rect, color=background_color) + + cv2.putText( + img=scene, + text=text, + org=(text_anchor.x - text_width // 2, text_anchor.y + text_height // 2), + fontFace=text_font, + fontScale=text_scale, + color=text_color.as_bgr(), + thickness=text_thickness, + lineType=cv2.LINE_AA, + ) + return scene From ccafa3339d88bb278689a2390a79ab3e2e4386e6 Mon Sep 17 00:00:00 2001 From: SkalskiP Date: Mon, 6 Feb 2023 23:05:17 +0100 Subject: [PATCH 09/23] - PolygonZone and PolygonZoneAnnotator added --- supervision/detection/core.py | 8 ++- supervision/detection/polygon_zone.py | 80 +++++++++++++++++++++++++++ supervision/draw/utils.py | 10 ++-- 3 files changed, 92 insertions(+), 6 deletions(-) diff --git a/supervision/detection/core.py b/supervision/detection/core.py index 2fef419b5..7e93001b8 100644 --- a/supervision/detection/core.py +++ b/supervision/detection/core.py @@ -150,9 +150,13 @@ def __getitem__(self, index: np.ndarray) -> Detections: xyxy=self.xyxy[index], confidence=self.confidence[index], class_id=self.class_id[index], - tracker_id=self.tracker_id[index] if self.tracker_id is not None else None, + tracker_id=self.tracker_id[index] + if self.tracker_id is not None + else None, ) - raise TypeError(f"Detections.__getitem__ not supported for index of type {type(index)}.") + raise TypeError( + f"Detections.__getitem__ not supported for index of type {type(index)}." + ) class BoxAnnotator: diff --git a/supervision/detection/polygon_zone.py b/supervision/detection/polygon_zone.py index e69de29bb..bc30a6b65 100644 --- a/supervision/detection/polygon_zone.py +++ b/supervision/detection/polygon_zone.py @@ -0,0 +1,80 @@ +from typing import Tuple + +import cv2 +import numpy as np + +from supervision import Detections +from supervision.detection.utils import generate_2d_mask +from supervision.draw.color import Color +from supervision.draw.utils import draw_polygon, draw_text +from supervision.geometry.core import Position +from supervision.geometry.utils import get_polygon_center + + +class PolygonZone: + def __init__( + self, + polygon: np.ndarray, + frame_resolution_wh: Tuple[int, int], + triggering_position: Position, + ): + self.polygon = polygon + self.frame_resolution_wh = frame_resolution_wh + self.triggering_position = triggering_position + self.mask = generate_2d_mask(polygon=polygon, resolution_wh=frame_resolution_wh) + self.current_count = 0 + + def trigger(self, detections: Detections) -> np.ndarray: + anchors = ( + np.ceil( + detections.get_anchor_coordinates(anchor=self.triggering_position) + ).astype(int) + - 1 + ) + is_in_zone = self.mask[anchors[:, 1], anchors[:, 0]] + self.current_count = np.sum(is_in_zone) + return is_in_zone.astype(bool) + + +class PolygonZoneAnnotator: + def __init__( + self, + zone: PolygonZone, + color: Color, + thickness: int = 2, + text_color: Color = Color.black(), + text_scale: float = 0.5, + text_thickness: int = 1, + text_padding: int = 10, + ): + self.zone = zone + self.color = color + self.thickness = thickness + self.text_color = text_color + self.text_scale = text_scale + self.text_thickness = text_thickness + self.text_padding = text_padding + self.font = cv2.FONT_HERSHEY_SIMPLEX + self.center = get_polygon_center(polygon=zone.polygon) + + def annotate(self, scene: np.ndarray) -> np.ndarray: + annotated_frame = draw_polygon( + scene=scene, + polygon=self.zone.polygon, + color=self.color, + thickness=self.thickness, + ) + + annotated_frame = draw_text( + scene=annotated_frame, + text=str(self.zone.current_count), + text_anchor=self.center, + background_color=self.color, + text_color=self.text_color, + text_scale=self.text_scale, + text_thickness=self.text_thickness, + text_padding=self.text_padding, + text_font=self.font, + ) + + return annotated_frame diff --git a/supervision/draw/utils.py b/supervision/draw/utils.py index ca9366de9..e8bcea212 100644 --- a/supervision/draw/utils.py +++ b/supervision/draw/utils.py @@ -1,11 +1,11 @@ +from typing import Optional + import cv2 import numpy as np from supervision.draw.color import Color from supervision.geometry.core import Point, Rect -from typing import Optional - def draw_line( scene: np.ndarray, start: Point, end: Point, color: Color, thickness: int = 2 @@ -152,11 +152,13 @@ def draw_text( x=text_anchor.x - text_width // 2, y=text_anchor.y - text_height // 2, width=text_width, - height=text_height + height=text_height, ).pad(text_padding) if background_color is not None: - scene = draw_filled_rectangle(scene=scene, rect=text_rect, color=background_color) + scene = draw_filled_rectangle( + scene=scene, rect=text_rect, color=background_color + ) cv2.putText( img=scene, From 8d091a7a004086ffe82f664400162270ba6a17ca Mon Sep 17 00:00:00 2001 From: SkalskiP Date: Tue, 7 Feb 2023 09:30:53 +0100 Subject: [PATCH 10/23] - Rename LineCounter to LineZone - Rename LineCounter.update() to LineZone.trigger() to make the API consistent with LineZone - Rename LineCounterAnnotator to LineZoneAnnotator --- supervision/__init__.py | 4 +++- supervision/detection/core.py | 2 +- supervision/detection/line_counter.py | 8 ++++---- 3 files changed, 8 insertions(+), 6 deletions(-) diff --git a/supervision/__init__.py b/supervision/__init__.py index 9a0f17def..2db5a72e8 100644 --- a/supervision/__init__.py +++ b/supervision/__init__.py @@ -1,3 +1,5 @@ __version__ = "0.2.0.dev0" -from supervision.detection.core import Detections +from supervision.video import get_video_frames_generator +from supervision.detection.core import Detections, BoxAnnotator +from supervision.notebook.utils import show_frame_in_notebook diff --git a/supervision/detection/core.py b/supervision/detection/core.py index 7e93001b8..155b9c346 100644 --- a/supervision/detection/core.py +++ b/supervision/detection/core.py @@ -66,7 +66,7 @@ def __iter__(self): ) @classmethod - def from_yolov5(cls, yolov5_detections: "yolov5.models.common.Detections"): + def from_yolov5(cls, yolov5_detections): """ Creates a Detections instance from a YOLOv5 output Detections diff --git a/supervision/detection/line_counter.py b/supervision/detection/line_counter.py index 030508cae..e40ead8ec 100644 --- a/supervision/detection/line_counter.py +++ b/supervision/detection/line_counter.py @@ -8,7 +8,7 @@ from supervision.geometry.core import Point, Rect, Vector -class LineCounter: +class LineZone: """ Count the number of objects that cross a line. """ @@ -27,7 +27,7 @@ def __init__(self, start: Point, end: Point): self.in_count: int = 0 self.out_count: int = 0 - def update(self, detections: Detections): + def trigger(self, detections: Detections): """ Update the in_count and out_count for the detections that cross the line. @@ -71,7 +71,7 @@ def update(self, detections: Detections): self.out_count += 1 -class LineCounterAnnotator: +class LineZoneAnnotator: def __init__( self, thickness: float = 2, @@ -103,7 +103,7 @@ def __init__( self.text_offset: float = text_offset self.text_padding: int = text_padding - def annotate(self, frame: np.ndarray, line_counter: LineCounter) -> np.ndarray: + def annotate(self, frame: np.ndarray, line_counter: LineZone) -> np.ndarray: """ Draws the line on the frame using the line_counter provided. From 4afc650eb2adda9b00696ce944052a310ff68533 Mon Sep 17 00:00:00 2001 From: SkalskiP Date: Tue, 7 Feb 2023 11:18:00 +0100 Subject: [PATCH 11/23] - skip_label argument added to BoxAnnotator.annotate() --- supervision/detection/core.py | 24 ++++++++++++++---------- 1 file changed, 14 insertions(+), 10 deletions(-) diff --git a/supervision/detection/core.py b/supervision/detection/core.py index 155b9c346..63ac7a912 100644 --- a/supervision/detection/core.py +++ b/supervision/detection/core.py @@ -193,20 +193,32 @@ def annotate( frame: np.ndarray, detections: Detections, labels: Optional[List[str]] = None, + skip_label: bool = False ) -> np.ndarray: """ Draws bounding boxes on the frame using the detections provided. - Attributes: + Parameters: frame (np.ndarray): The image on which the bounding boxes will be drawn detections (Detections): The detections for which the bounding boxes will be drawn labels (Optional[List[str]]): An optional list of labels corresponding to each detection. If labels is provided, the confidence score of the detection will be replaced with the label. - + skip_label (bool): Is set to True, skips bounding box label annotation. Returns: np.ndarray: The image with the bounding boxes drawn on it """ font = cv2.FONT_HERSHEY_SIMPLEX for i, (xyxy, confidence, class_id, tracker_id) in enumerate(detections): + x1, y1, x2, y2 = xyxy.astype(int) + cv2.rectangle( + img=frame, + pt1=(x1, y1), + pt2=(x2, y2), + color=color.as_bgr(), + thickness=self.thickness, + ) + if skip_label: + return frame + color = ( self.color.by_idx(class_id) if isinstance(self.color, ColorPalette) @@ -218,7 +230,6 @@ def annotate( else labels[i] ) - x1, y1, x2, y2 = xyxy.astype(int) text_width, text_height = cv2.getTextSize( text=text, fontFace=font, @@ -235,13 +246,6 @@ def annotate( text_background_x2 = x1 + 2 * self.text_padding + text_width text_background_y2 = y1 - cv2.rectangle( - img=frame, - pt1=(x1, y1), - pt2=(x2, y2), - color=color.as_bgr(), - thickness=self.thickness, - ) cv2.rectangle( img=frame, pt1=(text_background_x1, text_background_y1), From a64a252ae3377bee162d2346381e22138236cefc Mon Sep 17 00:00:00 2001 From: SkalskiP Date: Tue, 7 Feb 2023 11:31:37 +0100 Subject: [PATCH 12/23] - Detections.from_yolov8() added --- supervision/detection/core.py | 35 +++++++++++++++++++++++++++++++---- 1 file changed, 31 insertions(+), 4 deletions(-) diff --git a/supervision/detection/core.py b/supervision/detection/core.py index 63ac7a912..e7fc3cfc5 100644 --- a/supervision/detection/core.py +++ b/supervision/detection/core.py @@ -86,10 +86,37 @@ def from_yolov5(cls, yolov5_detections): ``` """ yolov5_detections_predictions = yolov5_detections.pred[0].cpu().cpu().numpy() - xyxy = yolov5_detections_predictions[:, :4] - confidence = yolov5_detections_predictions[:, 4] - class_id = yolov5_detections_predictions[:, 5].astype(int) - return cls(xyxy, confidence, class_id) + return cls( + xyxy=yolov5_detections_predictions[:, :4], + confidence=yolov5_detections_predictions[:, 4], + class_id=yolov5_detections_predictions[:, 5].astype(int) + ) + + @classmethod + def from_yolov8(cls, yolov8_results): + """ + Creates a Detections instance from a YOLOv8 output Results + + Attributes: + yolov8_results (ultralytics.yolo.engine.results.Results): The output Results instance from YOLOv8 + + Returns: + + Example: + ```python + >>> from ultralytics import YOLO + >>> from supervision import Detections + + >>> model = YOLO('yolov8s.pt') + >>> results = model(frame) + >>> detections = Detections.from_yolov8(results) + ``` + """ + 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) + ) def filter(self, mask: np.ndarray, inplace: bool = False) -> Optional[Detections]: """ From 940bcb8bc5d52e01b0a6ad290315e0e252992d4f Mon Sep 17 00:00:00 2001 From: SkalskiP Date: Tue, 7 Feb 2023 11:39:11 +0100 Subject: [PATCH 13/23] - BoxAnnotator.annotate() fixed --- supervision/detection/core.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/supervision/detection/core.py b/supervision/detection/core.py index e7fc3cfc5..5e63b68ce 100644 --- a/supervision/detection/core.py +++ b/supervision/detection/core.py @@ -236,6 +236,11 @@ def annotate( font = cv2.FONT_HERSHEY_SIMPLEX for i, (xyxy, confidence, class_id, tracker_id) in enumerate(detections): x1, y1, x2, y2 = xyxy.astype(int) + color = ( + self.color.by_idx(class_id) + if isinstance(self.color, ColorPalette) + else self.color + ) cv2.rectangle( img=frame, pt1=(x1, y1), @@ -246,11 +251,6 @@ def annotate( if skip_label: return frame - color = ( - self.color.by_idx(class_id) - if isinstance(self.color, ColorPalette) - else self.color - ) text = ( f"{confidence:0.2f}" if (labels is None or len(detections) != len(labels)) From 05705bd4f245b5ba41e492839078587d0f7f19f3 Mon Sep 17 00:00:00 2001 From: SkalskiP Date: Tue, 7 Feb 2023 11:49:16 +0100 Subject: [PATCH 14/23] - Fix linter issues --- supervision/__init__.py | 4 ++-- supervision/detection/core.py | 6 +++--- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/supervision/__init__.py b/supervision/__init__.py index 2db5a72e8..e3fa30cf7 100644 --- a/supervision/__init__.py +++ b/supervision/__init__.py @@ -1,5 +1,5 @@ __version__ = "0.2.0.dev0" -from supervision.video import get_video_frames_generator -from supervision.detection.core import Detections, BoxAnnotator +from supervision.detection.core import BoxAnnotator, Detections from supervision.notebook.utils import show_frame_in_notebook +from supervision.video import get_video_frames_generator diff --git a/supervision/detection/core.py b/supervision/detection/core.py index 5e63b68ce..754a656ac 100644 --- a/supervision/detection/core.py +++ b/supervision/detection/core.py @@ -89,7 +89,7 @@ def from_yolov5(cls, yolov5_detections): return cls( xyxy=yolov5_detections_predictions[:, :4], confidence=yolov5_detections_predictions[:, 4], - class_id=yolov5_detections_predictions[:, 5].astype(int) + class_id=yolov5_detections_predictions[:, 5].astype(int), ) @classmethod @@ -115,7 +115,7 @@ def from_yolov8(cls, yolov8_results): 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) + class_id=yolov8_results.boxes.cls.cpu().numpy().astype(int), ) def filter(self, mask: np.ndarray, inplace: bool = False) -> Optional[Detections]: @@ -220,7 +220,7 @@ def annotate( frame: np.ndarray, detections: Detections, labels: Optional[List[str]] = None, - skip_label: bool = False + skip_label: bool = False, ) -> np.ndarray: """ Draws bounding boxes on the frame using the detections provided. From 38bf4b0f544ec58ad1f6113c0d7b8bd484562cfb Mon Sep 17 00:00:00 2001 From: SkalskiP Date: Tue, 7 Feb 2023 11:52:07 +0100 Subject: [PATCH 15/23] - matplotlib added to setup.py --- setup.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/setup.py b/setup.py index 6afc4eaa5..d91926459 100644 --- a/setup.py +++ b/setup.py @@ -24,8 +24,9 @@ def get_version(): long_description_content_type='text/markdown', url='https://github.com/roboflow/supervision', install_requires=[ - 'numpy', - 'opencv-python' + 'numpy', + 'opencv-python', + 'matplotlib' ], packages=find_packages(exclude=("tests",)), extras_require={ From 51a1655bd45a03ca68c5705916b48a4a7a38e78a Mon Sep 17 00:00:00 2001 From: SkalskiP Date: Tue, 7 Feb 2023 12:06:56 +0100 Subject: [PATCH 16/23] - improve ease of import of supervision components --- supervision/__init__.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/supervision/__init__.py b/supervision/__init__.py index e3fa30cf7..bb6ee5b42 100644 --- a/supervision/__init__.py +++ b/supervision/__init__.py @@ -1,5 +1,11 @@ __version__ = "0.2.0.dev0" from supervision.detection.core import BoxAnnotator, Detections +from supervision.detection.polygon_zone import PolygonZone, PolygonZoneAnnotator +from supervision.detection.utils import generate_2d_mask +from supervision.draw.color import Color, ColorPalette +from supervision.draw.utils import draw_filled_rectangle, draw_polygon, draw_text +from supervision.geometry.core import Point, Position, Rect +from supervision.geometry.utils import get_polygon_center from supervision.notebook.utils import show_frame_in_notebook -from supervision.video import get_video_frames_generator +from supervision.video import VideoInfo, VideoSink, get_video_frames_generator From 52a8193b4aa8c416b7fd82652cf5d1be60274863 Mon Sep 17 00:00:00 2001 From: SkalskiP Date: Tue, 7 Feb 2023 12:24:56 +0100 Subject: [PATCH 17/23] - default value for triggering_position in PolygonZone added --- supervision/detection/polygon_zone.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/supervision/detection/polygon_zone.py b/supervision/detection/polygon_zone.py index bc30a6b65..af98c269f 100644 --- a/supervision/detection/polygon_zone.py +++ b/supervision/detection/polygon_zone.py @@ -16,7 +16,7 @@ def __init__( self, polygon: np.ndarray, frame_resolution_wh: Tuple[int, int], - triggering_position: Position, + triggering_position: Position = Position.BOTTOM_CENTER, ): self.polygon = polygon self.frame_resolution_wh = frame_resolution_wh From c8868c3412022844ffca0c835db88ff5c73a210f Mon Sep 17 00:00:00 2001 From: SkalskiP Date: Tue, 7 Feb 2023 12:34:59 +0100 Subject: [PATCH 18/23] - small API fixes in PolygonZoneAnnotator and BoxAnnotator --- supervision/detection/core.py | 14 +++++++------- supervision/detection/polygon_zone.py | 6 +++--- 2 files changed, 10 insertions(+), 10 deletions(-) diff --git a/supervision/detection/core.py b/supervision/detection/core.py index 754a656ac..ad5b6a72b 100644 --- a/supervision/detection/core.py +++ b/supervision/detection/core.py @@ -217,7 +217,7 @@ def __init__( def annotate( self, - frame: np.ndarray, + scene: np.ndarray, detections: Detections, labels: Optional[List[str]] = None, skip_label: bool = False, @@ -226,7 +226,7 @@ def annotate( Draws bounding boxes on the frame using the detections provided. Parameters: - frame (np.ndarray): The image on which the bounding boxes will be drawn + scene (np.ndarray): The image on which the bounding boxes will be drawn detections (Detections): The detections for which the bounding boxes will be drawn labels (Optional[List[str]]): An optional list of labels corresponding to each detection. If labels is provided, the confidence score of the detection will be replaced with the label. skip_label (bool): Is set to True, skips bounding box label annotation. @@ -242,14 +242,14 @@ def annotate( else self.color ) cv2.rectangle( - img=frame, + img=scene, pt1=(x1, y1), pt2=(x2, y2), color=color.as_bgr(), thickness=self.thickness, ) if skip_label: - return frame + return scene text = ( f"{confidence:0.2f}" @@ -274,14 +274,14 @@ def annotate( text_background_y2 = y1 cv2.rectangle( - img=frame, + img=scene, pt1=(text_background_x1, text_background_y1), pt2=(text_background_x2, text_background_y2), color=color.as_bgr(), thickness=cv2.FILLED, ) cv2.putText( - img=frame, + img=scene, text=text, org=(text_x, text_y), fontFace=font, @@ -290,4 +290,4 @@ def annotate( thickness=self.text_thickness, lineType=cv2.LINE_AA, ) - return frame + return scene diff --git a/supervision/detection/polygon_zone.py b/supervision/detection/polygon_zone.py index af98c269f..42601439e 100644 --- a/supervision/detection/polygon_zone.py +++ b/supervision/detection/polygon_zone.py @@ -1,4 +1,4 @@ -from typing import Tuple +from typing import Tuple, Optional import cv2 import numpy as np @@ -57,7 +57,7 @@ def __init__( self.font = cv2.FONT_HERSHEY_SIMPLEX self.center = get_polygon_center(polygon=zone.polygon) - def annotate(self, scene: np.ndarray) -> np.ndarray: + def annotate(self, scene: np.ndarray, label: Optional[str] = None) -> np.ndarray: annotated_frame = draw_polygon( scene=scene, polygon=self.zone.polygon, @@ -67,7 +67,7 @@ def annotate(self, scene: np.ndarray) -> np.ndarray: annotated_frame = draw_text( scene=annotated_frame, - text=str(self.zone.current_count), + text=str(self.zone.current_count) if label is None else label, text_anchor=self.center, background_color=self.color, text_color=self.text_color, From b8b08cacd3abbad4c6c883bb9151ec3a24bcaa97 Mon Sep 17 00:00:00 2001 From: SkalskiP Date: Tue, 7 Feb 2023 13:24:52 +0100 Subject: [PATCH 19/23] - make it simpler to import process_video function --- supervision/__init__.py | 2 +- supervision/video.py | 11 +++++------ 2 files changed, 6 insertions(+), 7 deletions(-) diff --git a/supervision/__init__.py b/supervision/__init__.py index bb6ee5b42..2d7228b42 100644 --- a/supervision/__init__.py +++ b/supervision/__init__.py @@ -8,4 +8,4 @@ from supervision.geometry.core import Point, Position, Rect from supervision.geometry.utils import get_polygon_center from supervision.notebook.utils import show_frame_in_notebook -from supervision.video import VideoInfo, VideoSink, get_video_frames_generator +from supervision.video import VideoInfo, VideoSink, get_video_frames_generator, process_video diff --git a/supervision/video.py b/supervision/video.py index 228b5a2d1..970f96cbb 100644 --- a/supervision/video.py +++ b/supervision/video.py @@ -20,7 +20,7 @@ class VideoInfo: Examples: ```python - >>> from supervision.video import VideoInfo + >>> from supervision import VideoInfo >>> video_info = VideoInfo.from_video_path(video_path='video.mp4') @@ -65,8 +65,7 @@ class VideoSink: Examples: ```python - >>> from supervision.video import VideoInfo - >>> from supervision.video import VideoSink + >>> from supervision import VideoInfo, VideoSink >>> video_info = VideoInfo.from_video_path(video_path='source_video.mp4') @@ -110,7 +109,7 @@ def get_video_frames_generator(source_path: str) -> Generator[np.ndarray, None, Examples: ```python - >>> from supervision.video import get_video_frames_generator + >>> from supervision import get_video_frames_generator >>> for frame in get_video_frames_generator(source_path='source_video.mp4'): ... ... @@ -141,9 +140,9 @@ def process_video( Examples: ```python - >>> from supervision.video import process_video + >>> from supervision import process_video - >>> def process_frame(frame: np.ndarray) -> np.ndarray: + >>> def process_frame(scene: np.ndarray) -> np.ndarray: ... ... >>> process_video( From 602d3af009dde8eec3c0ff969743ef67994084b5 Mon Sep 17 00:00:00 2001 From: SkalskiP Date: Tue, 7 Feb 2023 13:42:42 +0100 Subject: [PATCH 20/23] - fix bug in VideoSink --- supervision/video.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/supervision/video.py b/supervision/video.py index 970f96cbb..5dd807eb8 100644 --- a/supervision/video.py +++ b/supervision/video.py @@ -86,7 +86,7 @@ def __enter__(self): self.target_path, self.__fourcc, self.video_info.fps, - self.video_info.resolution, + self.video_info.resolution_wh, ) return self From 14d23eb9df1104f5b2ac69f57980e82357fa5fcf Mon Sep 17 00:00:00 2001 From: SkalskiP Date: Tue, 7 Feb 2023 14:14:04 +0100 Subject: [PATCH 21/23] - fix bug in BoxAnnotator --- supervision/detection/core.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/supervision/detection/core.py b/supervision/detection/core.py index ad5b6a72b..0863e64fc 100644 --- a/supervision/detection/core.py +++ b/supervision/detection/core.py @@ -249,7 +249,7 @@ def annotate( thickness=self.thickness, ) if skip_label: - return scene + continue text = ( f"{confidence:0.2f}" From fd35ad9653790f50f9b9cfc2bd6d18e4012ba010 Mon Sep 17 00:00:00 2001 From: SkalskiP Date: Tue, 7 Feb 2023 19:02:19 +0100 Subject: [PATCH 22/23] - set supervision version to 0.2.0 --- supervision/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/supervision/__init__.py b/supervision/__init__.py index 2d7228b42..4bf96cd08 100644 --- a/supervision/__init__.py +++ b/supervision/__init__.py @@ -1,4 +1,4 @@ -__version__ = "0.2.0.dev0" +__version__ = "0.2.0" from supervision.detection.core import BoxAnnotator, Detections from supervision.detection.polygon_zone import PolygonZone, PolygonZoneAnnotator From a0eb6c3237963dd6b53d2e3ff839578e0b944516 Mon Sep 17 00:00:00 2001 From: SkalskiP Date: Tue, 7 Feb 2023 19:03:25 +0100 Subject: [PATCH 23/23] - make linter happy --- supervision/__init__.py | 7 ++++++- supervision/detection/polygon_zone.py | 2 +- 2 files changed, 7 insertions(+), 2 deletions(-) diff --git a/supervision/__init__.py b/supervision/__init__.py index 4bf96cd08..c449f7733 100644 --- a/supervision/__init__.py +++ b/supervision/__init__.py @@ -8,4 +8,9 @@ from supervision.geometry.core import Point, Position, Rect from supervision.geometry.utils import get_polygon_center from supervision.notebook.utils import show_frame_in_notebook -from supervision.video import VideoInfo, VideoSink, get_video_frames_generator, process_video +from supervision.video import ( + VideoInfo, + VideoSink, + get_video_frames_generator, + process_video, +) diff --git a/supervision/detection/polygon_zone.py b/supervision/detection/polygon_zone.py index 42601439e..54bbbf625 100644 --- a/supervision/detection/polygon_zone.py +++ b/supervision/detection/polygon_zone.py @@ -1,4 +1,4 @@ -from typing import Tuple, Optional +from typing import Optional, Tuple import cv2 import numpy as np