From 550c913dbb6eef428ff600ec7670dac57986a458 Mon Sep 17 00:00:00 2001 From: SkalskiP Date: Thu, 23 Mar 2023 12:31:23 +0100 Subject: [PATCH 1/5] =?UTF-8?q?=F0=9F=92=A7=20drop=20requirement=20for=20c?= =?UTF-8?q?lass=5Fid=20in=20Detections=20constructor=20to=20make=20it=20mo?= =?UTF-8?q?re=20flexible?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- supervision/__init__.py | 5 +- supervision/detection/annotate.py | 114 ++++++++++++++++++++++++ supervision/detection/core.py | 140 +++++------------------------- 3 files changed, 139 insertions(+), 120 deletions(-) create mode 100644 supervision/detection/annotate.py diff --git a/supervision/__init__.py b/supervision/__init__.py index c198ee574..0980d7c14 100644 --- a/supervision/__init__.py +++ b/supervision/__init__.py @@ -1,6 +1,7 @@ -__version__ = "0.3.1" +__version__ = "0.3.2" -from supervision.detection.core import BoxAnnotator, Detections +from supervision.detection.core import Detections +from supervision.detection.annotate import BoxAnnotator from supervision.detection.line_counter import LineZone, LineZoneAnnotator from supervision.detection.polygon_zone import PolygonZone, PolygonZoneAnnotator from supervision.detection.utils import generate_2d_mask diff --git a/supervision/detection/annotate.py b/supervision/detection/annotate.py new file mode 100644 index 000000000..de70ec4e4 --- /dev/null +++ b/supervision/detection/annotate.py @@ -0,0 +1,114 @@ +from typing import Union, Optional, List + +import cv2 +import numpy as np + +from supervision import Color, ColorPalette, Detections + + +class BoxAnnotator: + def __init__( + self, + color: Union[Color, ColorPalette] = ColorPalette.default(), + thickness: int = 2, + text_color: Color = Color.black(), + text_scale: float = 0.5, + text_thickness: int = 1, + text_padding: int = 10, + ): + """ + A class for drawing bounding boxes on an image using detections provided. + + Attributes: + color (Union[Color, ColorPalette]): The color to draw the bounding box, can be a single color or a color palette + thickness (int): The thickness of the bounding box lines, default is 2 + text_color (Color): The color of the text on the bounding box, default is white + text_scale (float): The scale of the text on the bounding box, default is 0.5 + text_thickness (int): The thickness of the text on the bounding box, default is 1 + text_padding (int): The padding around the text on the bounding box, default is 5 + + """ + self.color: Union[Color, ColorPalette] = color + self.thickness: int = thickness + self.text_color: Color = text_color + self.text_scale: float = text_scale + self.text_thickness: int = text_thickness + self.text_padding: int = text_padding + + def annotate( + self, + scene: 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. + + Parameters: + 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. + 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) + idx = class_id if not None else i + color = ( + self.color.by_idx(idx) + if isinstance(self.color, ColorPalette) + else self.color + ) + cv2.rectangle( + img=scene, + pt1=(x1, y1), + pt2=(x2, y2), + color=color.as_bgr(), + thickness=self.thickness, + ) + if skip_label: + continue + + text = ( + f"{class_id}" + if (labels is None or len(detections) != len(labels)) + else labels[i] + ) + + text_width, text_height = cv2.getTextSize( + text=text, + fontFace=font, + fontScale=self.text_scale, + thickness=self.text_thickness, + )[0] + + text_x = x1 + self.text_padding + text_y = y1 - self.text_padding + + text_background_x1 = x1 + text_background_y1 = y1 - 2 * self.text_padding - text_height + + text_background_x2 = x1 + 2 * self.text_padding + text_width + text_background_y2 = y1 + + cv2.rectangle( + 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=scene, + text=text, + org=(text_x, text_y), + fontFace=font, + fontScale=self.text_scale, + color=self.text_color.as_rgb(), + thickness=self.text_thickness, + lineType=cv2.LINE_AA, + ) + return scene diff --git a/supervision/detection/core.py b/supervision/detection/core.py index a2da1b3f8..bca273de0 100644 --- a/supervision/detection/core.py +++ b/supervision/detection/core.py @@ -1,13 +1,11 @@ from __future__ import annotations from dataclasses import dataclass -from typing import Iterator, List, Optional, Tuple, Union +from typing import Iterator, Optional, Tuple, Union -import cv2 import numpy as np from supervision.detection.utils import non_max_suppression -from supervision.draw.color import Color, ColorPalette from supervision.geometry.core import Position @@ -18,13 +16,13 @@ class Detections: Attributes: xyxy (np.ndarray): An array of shape `(n, 4)` containing the bounding boxes coordinates in format `[x1, y1, x2, y2]` + class_id (Optional[np.ndarray]): An array of shape `(n,)` containing the class ids of the detections. confidence (Optional[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 - class_id: np.ndarray + class_id: Optional[np.ndarray] = None confidence: Optional[np.ndarray] = None tracker_id: Optional[np.ndarray] = None @@ -32,7 +30,11 @@ def __post_init__(self): n = len(self.xyxy) validators = [ (isinstance(self.xyxy, np.ndarray) and self.xyxy.shape == (n, 4)), - (isinstance(self.class_id, np.ndarray) and self.class_id.shape == (n,)), + self.class_id is None + or ( + isinstance(self.class_id, np.ndarray) + and self.class_id.shape == (n,) + ), self.confidence is None or ( isinstance(self.confidence, np.ndarray) @@ -47,8 +49,8 @@ def __post_init__(self): if not all(validators): raise ValueError( "xyxy must be 2d np.ndarray with (n, 4) shape, " + "class_id must be None or 1d np.ndarray with (n,) shape, " "confidence must be None or 1d np.ndarray with (n,) shape, " - "class_id must be 1d np.ndarray with (n,) shape, " "tracker_id must be None or 1d np.ndarray with (n,) shape" ) @@ -68,7 +70,7 @@ def __iter__( yield ( self.xyxy[i], self.confidence[i] if self.confidence is not None else None, - self.class_id[i], + self.class_id[i] if self.class_id is not None else None, self.tracker_id[i] if self.tracker_id is not None else None, ) @@ -76,13 +78,18 @@ def __eq__(self, other: Detections): return all( [ np.array_equal(self.xyxy, other.xyxy), + any( + [ + self.class_id is None and other.class_id is None, + np.array_equal(self.class_id, other.class_id), + ] + ), any( [ self.confidence is None and other.confidence is None, np.array_equal(self.confidence, other.confidence), ] ), - np.array_equal(self.class_id, other.class_id), any( [ self.tracker_id is None and other.tracker_id is None, @@ -213,8 +220,12 @@ def __getitem__(self, index: np.ndarray) -> Detections: ): return Detections( xyxy=self.xyxy[index], - confidence=self.confidence[index], - class_id=self.class_id[index], + confidence=self.confidence[index] + if self.confidence is not None + else None, + class_id=self.class_id[index] + if self.class_id is not None + else None, tracker_id=self.tracker_id[index] if self.tracker_id is not None else None, @@ -273,110 +284,3 @@ def with_nms( ) indices = non_max_suppression(predictions=predictions, iou_threshold=threshold) return self[indices] - - -class BoxAnnotator: - def __init__( - self, - color: Union[Color, ColorPalette] = ColorPalette.default(), - thickness: int = 2, - text_color: Color = Color.black(), - text_scale: float = 0.5, - text_thickness: int = 1, - text_padding: int = 10, - ): - """ - A class for drawing bounding boxes on an image using detections provided. - - Attributes: - color (Union[Color, ColorPalette]): The color to draw the bounding box, can be a single color or a color palette - thickness (int): The thickness of the bounding box lines, default is 2 - text_color (Color): The color of the text on the bounding box, default is white - text_scale (float): The scale of the text on the bounding box, default is 0.5 - text_thickness (int): The thickness of the text on the bounding box, default is 1 - text_padding (int): The padding around the text on the bounding box, default is 5 - - """ - self.color: Union[Color, ColorPalette] = color - self.thickness: int = thickness - self.text_color: Color = text_color - self.text_scale: float = text_scale - self.text_thickness: int = text_thickness - self.text_padding: int = text_padding - - def annotate( - self, - scene: 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. - - Parameters: - 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. - 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) - color = ( - self.color.by_idx(class_id) - if isinstance(self.color, ColorPalette) - else self.color - ) - cv2.rectangle( - img=scene, - pt1=(x1, y1), - pt2=(x2, y2), - color=color.as_bgr(), - thickness=self.thickness, - ) - if skip_label: - continue - - text = ( - f"{class_id}" - if (labels is None or len(detections) != len(labels)) - else labels[i] - ) - - text_width, text_height = cv2.getTextSize( - text=text, - fontFace=font, - fontScale=self.text_scale, - thickness=self.text_thickness, - )[0] - - text_x = x1 + self.text_padding - text_y = y1 - self.text_padding - - text_background_x1 = x1 - text_background_y1 = y1 - 2 * self.text_padding - text_height - - text_background_x2 = x1 + 2 * self.text_padding + text_width - text_background_y2 = y1 - - cv2.rectangle( - 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=scene, - text=text, - org=(text_x, text_y), - fontFace=font, - fontScale=self.text_scale, - color=self.text_color.as_rgb(), - thickness=self.text_thickness, - lineType=cv2.LINE_AA, - ) - return scene From 420dae9e70952508717dbd3b76a8ba1f3f56024d Mon Sep 17 00:00:00 2001 From: SkalskiP Date: Thu, 23 Mar 2023 13:02:52 +0100 Subject: [PATCH 2/5] =?UTF-8?q?=F0=9F=9B=A0=EF=B8=8F=20fix=20circular=20im?= =?UTF-8?q?port?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- supervision/detection/annotate.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/supervision/detection/annotate.py b/supervision/detection/annotate.py index de70ec4e4..de6e314fb 100644 --- a/supervision/detection/annotate.py +++ b/supervision/detection/annotate.py @@ -3,7 +3,8 @@ import cv2 import numpy as np -from supervision import Color, ColorPalette, Detections +from supervision.draw.color import Color, ColorPalette +from supervision.detection.core import Detections class BoxAnnotator: From b2f6520b87ec3bf046e1d4fcbd651a497e5dbfe2 Mon Sep 17 00:00:00 2001 From: SkalskiP Date: Thu, 23 Mar 2023 13:43:03 +0100 Subject: [PATCH 3/5] =?UTF-8?q?=F0=9F=9B=A0=EF=B8=8F=20fix=20idx=20error?= =?UTF-8?q?=20in=20annotate?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- supervision/detection/annotate.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/supervision/detection/annotate.py b/supervision/detection/annotate.py index de6e314fb..f8bfff553 100644 --- a/supervision/detection/annotate.py +++ b/supervision/detection/annotate.py @@ -57,7 +57,7 @@ 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) - idx = class_id if not None else i + idx = class_id if class_id is not None else i color = ( self.color.by_idx(idx) if isinstance(self.color, ColorPalette) From 8bf686e23beb4e3283b489d185717611766b13fd Mon Sep 17 00:00:00 2001 From: SkalskiP Date: Thu, 23 Mar 2023 14:18:46 +0100 Subject: [PATCH 4/5] =?UTF-8?q?=F0=9F=96=A4=20make=20black=20happy?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- mkdocs.yml | 4 ++-- supervision/__init__.py | 2 +- supervision/detection/annotate.py | 4 ++-- supervision/detection/core.py | 9 ++------- 4 files changed, 7 insertions(+), 12 deletions(-) diff --git a/mkdocs.yml b/mkdocs.yml index eaf8339a6..9c2672758 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -34,8 +34,8 @@ nav: theme: name: 'material' - logo: https://raw.githubusercontent.com/roboflow/supervision/main/docs/assets/roboflow_logomark_white.svg - favicon: https://raw.githubusercontent.com/roboflow/supervision/main/docs/assets/roboflow_logomark_color.svg + logo: https://media.roboflow.com/open-source/supervision/supervision-lenny.png?updatedAt=1678995918671 + favicon: https://media.roboflow.com/open-source/supervision/supervision-lenny.png?updatedAt=1678995918671 palette: primary: 'deep purple' accent: 'teal' diff --git a/supervision/__init__.py b/supervision/__init__.py index 0980d7c14..62c309d08 100644 --- a/supervision/__init__.py +++ b/supervision/__init__.py @@ -1,7 +1,7 @@ __version__ = "0.3.2" -from supervision.detection.core import Detections from supervision.detection.annotate import BoxAnnotator +from supervision.detection.core import Detections from supervision.detection.line_counter import LineZone, LineZoneAnnotator from supervision.detection.polygon_zone import PolygonZone, PolygonZoneAnnotator from supervision.detection.utils import generate_2d_mask diff --git a/supervision/detection/annotate.py b/supervision/detection/annotate.py index f8bfff553..29f0fba58 100644 --- a/supervision/detection/annotate.py +++ b/supervision/detection/annotate.py @@ -1,10 +1,10 @@ -from typing import Union, Optional, List +from typing import List, Optional, Union import cv2 import numpy as np -from supervision.draw.color import Color, ColorPalette from supervision.detection.core import Detections +from supervision.draw.color import Color, ColorPalette class BoxAnnotator: diff --git a/supervision/detection/core.py b/supervision/detection/core.py index bca273de0..431ec3f16 100644 --- a/supervision/detection/core.py +++ b/supervision/detection/core.py @@ -31,10 +31,7 @@ def __post_init__(self): validators = [ (isinstance(self.xyxy, np.ndarray) and self.xyxy.shape == (n, 4)), self.class_id is None - or ( - isinstance(self.class_id, np.ndarray) - and self.class_id.shape == (n,) - ), + or (isinstance(self.class_id, np.ndarray) and self.class_id.shape == (n,)), self.confidence is None or ( isinstance(self.confidence, np.ndarray) @@ -223,9 +220,7 @@ def __getitem__(self, index: np.ndarray) -> Detections: confidence=self.confidence[index] if self.confidence is not None else None, - class_id=self.class_id[index] - if self.class_id is not None - else None, + class_id=self.class_id[index] if self.class_id is not None else None, tracker_id=self.tracker_id[index] if self.tracker_id is not None else None, From a88eff868e847da7b9e513e2726352e7f05f91c7 Mon Sep 17 00:00:00 2001 From: SkalskiP Date: Thu, 23 Mar 2023 14:41:03 +0100 Subject: [PATCH 5/5] =?UTF-8?q?=F0=9F=93=96=20improve=20documentation?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- docs/detection/annotate.md | 3 +++ docs/{detection_core.md => detection/core.md} | 0 .../utils.md} | 0 mkdocs.yml | 5 ++-- supervision/detection/annotate.py | 25 ++++++++++--------- 5 files changed, 19 insertions(+), 14 deletions(-) create mode 100644 docs/detection/annotate.md rename docs/{detection_core.md => detection/core.md} (100%) rename docs/{detection_utils.md => detection/utils.md} (100%) diff --git a/docs/detection/annotate.md b/docs/detection/annotate.md new file mode 100644 index 000000000..fcebdec70 --- /dev/null +++ b/docs/detection/annotate.md @@ -0,0 +1,3 @@ +## BoxAnnotator + +:::supervision.detection.annotate.BoxAnnotator \ No newline at end of file diff --git a/docs/detection_core.md b/docs/detection/core.md similarity index 100% rename from docs/detection_core.md rename to docs/detection/core.md diff --git a/docs/detection_utils.md b/docs/detection/utils.md similarity index 100% rename from docs/detection_utils.md rename to docs/detection/utils.md diff --git a/mkdocs.yml b/mkdocs.yml index 9c2672758..254e775b5 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -26,8 +26,9 @@ nav: - API reference: - Video: video.md - Detection: - - Core: detection_core.md - - Utils: detection_utils.md + - Core: detection/core.md + - Annotate: detection/annotate.md + - Utils: detection/utils.md - Draw: - Utils: draw_utils.md - Notebook: notebook.md diff --git a/supervision/detection/annotate.py b/supervision/detection/annotate.py index 29f0fba58..3c74653d0 100644 --- a/supervision/detection/annotate.py +++ b/supervision/detection/annotate.py @@ -8,6 +8,19 @@ class BoxAnnotator: + """ + A class for drawing bounding boxes on an image using detections provided. + + Attributes: + color (Union[Color, ColorPalette]): The color to draw the bounding box, can be a single color or a color palette + thickness (int): The thickness of the bounding box lines, default is 2 + text_color (Color): The color of the text on the bounding box, default is white + text_scale (float): The scale of the text on the bounding box, default is 0.5 + text_thickness (int): The thickness of the text on the bounding box, default is 1 + text_padding (int): The padding around the text on the bounding box, default is 5 + + """ + def __init__( self, color: Union[Color, ColorPalette] = ColorPalette.default(), @@ -17,18 +30,6 @@ def __init__( text_thickness: int = 1, text_padding: int = 10, ): - """ - A class for drawing bounding boxes on an image using detections provided. - - Attributes: - color (Union[Color, ColorPalette]): The color to draw the bounding box, can be a single color or a color palette - thickness (int): The thickness of the bounding box lines, default is 2 - text_color (Color): The color of the text on the bounding box, default is white - text_scale (float): The scale of the text on the bounding box, default is 0.5 - text_thickness (int): The thickness of the text on the bounding box, default is 1 - text_padding (int): The padding around the text on the bounding box, default is 5 - - """ self.color: Union[Color, ColorPalette] = color self.thickness: int = thickness self.text_color: Color = text_color