From 9b77b2a1c6f32432181bed7731faec687ac8f0cb Mon Sep 17 00:00:00 2001 From: hd Date: Fri, 30 Jun 2023 18:25:13 +0200 Subject: [PATCH 01/17] initial commit for advanced annotator --- supervision/detection/annotate.py | 173 ++++++++++++++++++++++++++++++ 1 file changed, 173 insertions(+) diff --git a/supervision/detection/annotate.py b/supervision/detection/annotate.py index 3b3e40da6..17c78df32 100644 --- a/supervision/detection/annotate.py +++ b/supervision/detection/annotate.py @@ -37,6 +37,179 @@ def __init__( self.text_thickness: int = text_thickness self.text_padding: int = text_padding + def annotate_advance( + self, + scene: np.ndarray, + detections: Detections, + labels: Optional[List[str]] = None, + skip_label: bool = False, + ): + """ + Draws bounding boxes on the frame using the detections provided. + + Args: + 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` are not provided, corresponding `class_id` will be used as 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 + + Example: + ```python + >>> import supervision as sv + + >>> classes = ['person', ...] + >>> image = ... + >>> detections = sv.Detections(...) + + >>> box_annotator = sv.BoxAnnotator() + >>> labels = [ + ... f"{classes[class_id]} {confidence:0.2f}" + ... for _, _, confidence, class_id, _ + ... in detections + ... ] + >>> annotated_frame = box_annotator.annotate( + ... scene=image.copy(), + ... detections=detections, + ... labels=labels + ... ) + ``` + """ + font = cv2.FONT_HERSHEY_SIMPLEX + overlay_img = np.zeros_like(scene, np.uint8) + alpha = 0.8 + + line_thickness = self.thickness + 2 + for i in range(len(detections)): + x1, y1, x2, y2 = detections.xyxy[i].astype(int) + class_id = ( + detections.class_id[i] if detections.class_id is not None else None + ) + idx = class_id if class_id is not None else i + color = ( + self.color.by_idx(idx) + if isinstance(self.color, ColorPalette) + else self.color + ) + + box_width = x2 - x1 + box_height = y2 - y1 + + cv2.line( + scene, + (x1, y1), + (x1 + int(0.2 * box_width), y1), + color.as_bgr(), + thickness=line_thickness, + ) + cv2.line( + scene, + (x2 - int(0.2 * box_width), y1), + (x2, y1), + color.as_bgr(), + thickness=line_thickness, + ) + + cv2.line( + scene, + (x1, y2), + (x1 + int(0.2 * box_width), y2), + color.as_bgr(), + thickness=line_thickness, + ) + cv2.line( + scene, + (x2 - int(0.2 * box_width), y2), + (x2, y2), + color.as_bgr(), + thickness=line_thickness, + ) + # + cv2.line( + scene, + (x1, y1), + (x1, y1 + int(0.2 * box_height)), + color.as_bgr(), + thickness=line_thickness, + ) + cv2.line( + scene, + (x2, y1 + int(0.2 * box_height)), + (x2, y1), + color.as_bgr(), + thickness=line_thickness, + ) + + cv2.line( + scene, + (x1, y2 - int(0.2 * box_height)), + (x1, y2), + color.as_bgr(), + thickness=line_thickness, + ) + cv2.line( + scene, + (x2, y2 - int(0.2 * box_height)), + (x2, y2), + color.as_bgr(), + thickness=line_thickness, + ) + + cv2.rectangle( + img=overlay_img, + pt1=(x1, y1), + pt2=(x2, y2), + color=color.as_bgr(), + thickness=-1, + ) + 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, + ) + + mask = overlay_img.astype(bool) + scene[mask] = cv2.addWeighted(scene, alpha, overlay_img, 1 - alpha, 0)[mask] + return scene + def annotate( self, scene: np.ndarray, From b5e13420384d1973bccfd3358544590f96a08e31 Mon Sep 17 00:00:00 2001 From: hd Date: Tue, 4 Jul 2023 13:38:10 +0200 Subject: [PATCH 02/17] Updating refactored annotators --- supervision/__init__.py | 4 +- supervision/annotators/composable.py | 57 +++++++- supervision/annotators/core.py | 197 +++++++++++++++++++++++++-- 3 files changed, 245 insertions(+), 13 deletions(-) diff --git a/supervision/__init__.py b/supervision/__init__.py index 30694450b..cbc5de396 100644 --- a/supervision/__init__.py +++ b/supervision/__init__.py @@ -6,7 +6,9 @@ ClassificationDataset, DetectionDataset, ) -from supervision.detection.annotate import BoxAnnotator, MaskAnnotator + +from supervision.annotators.core import LabelAnnotator, MaskAnnotator, BoxAnnotator, PillowLabelAnnotator +from supervision.annotators.composable import DetectionAnnotator, TrackAnnotator, SegmentationAnnotator from supervision.detection.core import Detections from supervision.detection.line_counter import LineZone, LineZoneAnnotator from supervision.detection.tools.polygon_zone import PolygonZone, PolygonZoneAnnotator diff --git a/supervision/annotators/composable.py b/supervision/annotators/composable.py index 4c0eaeb98..a5f0c34c1 100644 --- a/supervision/annotators/composable.py +++ b/supervision/annotators/composable.py @@ -1,9 +1,9 @@ from abc import ABC -from typing import List +from typing import List, Optional import numpy as np -from supervision.annotators.core import BaseAnnotator +from supervision.annotators.core import BaseAnnotator, BoxAnnotator, LabelAnnotator, MaskAnnotator, TrackAnnotator from supervision.detection.core import Detections @@ -16,3 +16,56 @@ def annotate(self, scene: np.ndarray, detections: Detections) -> np.ndarray: for annotator in self.annotators: annotated_image = annotator.annotate(scene=scene, detections=detections) return annotated_image + + +class DetectionAnnotator(ComposableAnnotator): + def __init__(self): + self.annotators = [ + BoxAnnotator(), + LabelAnnotator() + ] + + def annotate(self, scene: np.ndarray, detections: Detections, labels: Optional[List[str]] = None,) -> np.ndarray: + for annotator in self.annotators: + if isinstance(annotator, LabelAnnotator): + scene = annotator.annotate(scene=scene, detections=detections, labels=labels) + else: + scene = annotator.annotate(scene=scene, detections=detections) + return scene + + +class SegmentationAnnotator(ComposableAnnotator): + def __init__(self): + self.annotators = [ + BoxAnnotator(), + LabelAnnotator(), + MaskAnnotator(), + ] + + def annotate(self, scene: np.ndarray, detections: Detections, labels: Optional[List[str]] = None,) -> np.ndarray: + for annotator in self.annotators: + if isinstance(annotator, LabelAnnotator): + scene = annotator.annotate(scene=scene, detections=detections, labels=labels) + else: + scene = annotator.annotate(scene=scene, detections=detections) + return scene + + +class TrackedDetectionAnnotator(ComposableAnnotator): + + def __init__(self): + self.annotators = [ + BoxAnnotator(), + LabelAnnotator(), + TrackAnnotator() + ] + + def annotate(self, scene: np.ndarray, detections: Detections, + labels: Optional[List[str]] = None, + color_by_track: bool = False,) -> np.ndarray: + for annotator in self.annotators: + if isinstance(annotator, LabelAnnotator): + scene = annotator.annotate(scene=scene, detections=detections, labels=labels) + else: + scene = annotator.annotate(scene=scene, detections=detections) + return scene \ No newline at end of file diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index 9e477a1aa..a81ce43fb 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -1,7 +1,8 @@ from abc import ABC, abstractmethod -from typing import Union +from typing import Union, Optional, Dict, List import numpy as np +import cv2 from supervision.detection.core import Detections from supervision.draw.color import Color, ColorPalette @@ -14,8 +15,72 @@ def annotate(self, scene: np.ndarray, detections: Detections) -> np.ndarray: class BoxAnnotator(BaseAnnotator): - def annotate(self, scene: np.ndarray, detections: Detections) -> np.ndarray: - pass + def __init__( + self, + color: Union[Color, ColorPalette] = ColorPalette.default(), + thickness: int = 2, + ): + self.color: Union[Color, ColorPalette] = color + self.thickness: int = thickness + + def annotate( + self, + scene: np.ndarray, + detections: Detections, + color_by_track: bool = False, + ) -> np.ndarray: + """ + Draws bounding boxes on the frame using the detections provided. + + Args: + 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 + color_by_track (bool): It allows to pick color by tracker id if present + Returns: + np.ndarray: The image with the bounding boxes drawn on it + + Example: + ```python + >>> import supervision as sv + + >>> classes = ['person', ...] + >>> image = ... + >>> detections = sv.Detections(...) + + >>> box_annotator = sv.BoxAnnotator() + >>> annotated_frame = box_annotator.annotate( + ... scene=image.copy(), + ... detections=detections + ... ) + ``` + """ + for i in range(len(detections)): + x1, y1, x2, y2 = detections.xyxy[i].astype(int) + + if color_by_track: + tracker_id = ( + detections.tracker_id[i] if detections.tracker_id is not None else None + ) + idx = tracker_id if tracker_id is not None else i + else: + class_id = ( + detections.class_id[i] if detections.class_id is not None else None + ) + idx = class_id if class_id is 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, + ) + + return scene class MaskAnnotator(BaseAnnotator): @@ -35,7 +100,7 @@ def __init__( self.color: Union[Color, ColorPalette] = color self.opacity = opacity - def annotate(self, scene: np.ndarray, detections: Detections) -> np.ndarray: + def annotate(self, scene: np.ndarray, detections: Detections, color_by_track: bool = False,) -> np.ndarray: """ Overlays the masks on the given image based on the provided detections, with a specified opacity. @@ -50,10 +115,17 @@ def annotate(self, scene: np.ndarray, detections: Detections) -> np.ndarray: return scene for i in np.flip(np.argsort(detections.area)): - class_id = ( - detections.class_id[i] if detections.class_id is not None else None - ) - idx = class_id if class_id is not None else i + if color_by_track: + tracker_id = ( + detections.tracker_id[i] if detections.tracker_id is not None else None + ) + idx = tracker_id if tracker_id is not None else i + else: + class_id = ( + detections.class_id[i] if detections.class_id is not None else None + ) + idx = class_id if class_id is not None else i + color = ( self.color.by_idx(idx) if isinstance(self.color, ColorPalette) @@ -74,8 +146,113 @@ def annotate(self, scene: np.ndarray, detections: Detections) -> np.ndarray: class LabelAnnotator(BaseAnnotator): - def annotate(self, scene: np.ndarray, detections: Detections) -> np.ndarray: - pass + + 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,): + 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, + color_by_track: bool = False,) -> np.ndarray: + """ + Draws bounding boxes on the frame using the detections provided. + + Args: + 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` are not provided, corresponding `class_id` will be used as 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 + + Example: + ```python + >>> import supervision as sv + + >>> classes = ['person', ...] + >>> image = ... + >>> detections = sv.Detections(...) + + >>> label_annotator = sv.LabelAnnotator() + >>> labels = [ + ... f"{classes[class_id]} {confidence:0.2f}" + ... for _, _, confidence, class_id, _ + ... in detections + ... ] + >>> annotated_frame = label_annotator.annotate( + ... scene=image.copy(), + ... detections=detections, + ... labels=labels + ... ) + ``` + """ + font = cv2.FONT_HERSHEY_SIMPLEX + for i in range(len(detections)): + x1, y1, x2, y2 = detections.xyxy[i].astype(int) + class_id = ( + detections.class_id[i] if detections.class_id is not None else None + ) + idx = class_id if class_id is not None else i + color = ( + self.color.by_idx(idx) + if isinstance(self.color, ColorPalette) + else self.color + ) + + 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 class PillowLabelAnnotator(BaseAnnotator): From 1bcf8346dc51046cc8179bafcaeca1834f51b1d6 Mon Sep 17 00:00:00 2001 From: hd Date: Tue, 4 Jul 2023 16:12:30 +0200 Subject: [PATCH 03/17] updated code --- supervision/__init__.py | 2 +- supervision/annotators/composable.py | 60 +++++++++++----------------- supervision/annotators/core.py | 10 +++++ 3 files changed, 34 insertions(+), 38 deletions(-) diff --git a/supervision/__init__.py b/supervision/__init__.py index cbc5de396..5ccabe3e6 100644 --- a/supervision/__init__.py +++ b/supervision/__init__.py @@ -8,7 +8,7 @@ ) from supervision.annotators.core import LabelAnnotator, MaskAnnotator, BoxAnnotator, PillowLabelAnnotator -from supervision.annotators.composable import DetectionAnnotator, TrackAnnotator, SegmentationAnnotator +from supervision.annotators.composable import DetectionAnnotator, TrackAnnotator, SegmentationAnnotator, TrackedDetectionAnnotator from supervision.detection.core import Detections from supervision.detection.line_counter import LineZone, LineZoneAnnotator from supervision.detection.tools.polygon_zone import PolygonZone, PolygonZoneAnnotator diff --git a/supervision/annotators/composable.py b/supervision/annotators/composable.py index a5f0c34c1..58b6a2aae 100644 --- a/supervision/annotators/composable.py +++ b/supervision/annotators/composable.py @@ -3,69 +3,55 @@ import numpy as np -from supervision.annotators.core import BaseAnnotator, BoxAnnotator, LabelAnnotator, MaskAnnotator, TrackAnnotator +from supervision.annotators.core import BoxAnnotator, LabelAnnotator, MaskAnnotator, TrackAnnotator from supervision.detection.core import Detections class ComposableAnnotator(ABC): - def __init__(self, annotators: List[BaseAnnotator]): - self.annotators = annotators + def __init__(self, ): + self.annotators = [] - def annotate(self, scene: np.ndarray, detections: Detections) -> np.ndarray: - annotated_image = scene + def annotate(self, scene: np.ndarray, detections: Detections, labels: Optional[List[str]] = None, + color_by_track: bool = False,) -> np.ndarray: for annotator in self.annotators: - annotated_image = annotator.annotate(scene=scene, detections=detections) - return annotated_image + if isinstance(annotator, LabelAnnotator): + scene = annotator.annotate(scene=scene, detections=detections, labels=labels, color_by_track=color_by_track) + else: + scene = annotator.annotate(scene=scene, detections=detections, color_by_track=color_by_track) + return scene class DetectionAnnotator(ComposableAnnotator): def __init__(self): + super().__init__() self.annotators = [ BoxAnnotator(), - LabelAnnotator() + LabelAnnotator(), ] - def annotate(self, scene: np.ndarray, detections: Detections, labels: Optional[List[str]] = None,) -> np.ndarray: - for annotator in self.annotators: - if isinstance(annotator, LabelAnnotator): - scene = annotator.annotate(scene=scene, detections=detections, labels=labels) - else: - scene = annotator.annotate(scene=scene, detections=detections) - return scene - class SegmentationAnnotator(ComposableAnnotator): + """ + A class for drawing segmentation mask, bounding box and labels on an image using provided detections. + """ def __init__(self): + super().__init__() self.annotators = [ BoxAnnotator(), LabelAnnotator(), MaskAnnotator(), ] - - def annotate(self, scene: np.ndarray, detections: Detections, labels: Optional[List[str]] = None,) -> np.ndarray: - for annotator in self.annotators: - if isinstance(annotator, LabelAnnotator): - scene = annotator.annotate(scene=scene, detections=detections, labels=labels) - else: - scene = annotator.annotate(scene=scene, detections=detections) - return scene + raise NotImplementedError class TrackedDetectionAnnotator(ComposableAnnotator): - + """ + A class for drawing object trajectories, bounding box and tracker ids on an image using provided detections. + """ def __init__(self): + super().__init__() self.annotators = [ BoxAnnotator(), LabelAnnotator(), - TrackAnnotator() - ] - - def annotate(self, scene: np.ndarray, detections: Detections, - labels: Optional[List[str]] = None, - color_by_track: bool = False,) -> np.ndarray: - for annotator in self.annotators: - if isinstance(annotator, LabelAnnotator): - scene = annotator.annotate(scene=scene, detections=detections, labels=labels) - else: - scene = annotator.annotate(scene=scene, detections=detections) - return scene \ No newline at end of file + TrackAnnotator(), + ] \ No newline at end of file diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index a81ce43fb..47a14aeed 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -15,6 +15,9 @@ def annotate(self, scene: np.ndarray, detections: Detections) -> np.ndarray: class BoxAnnotator(BaseAnnotator): + """ + Basic bounding box annotation class + """ def __init__( self, color: Union[Color, ColorPalette] = ColorPalette.default(), @@ -146,6 +149,12 @@ def annotate(self, scene: np.ndarray, detections: Detections, color_by_track: bo class LabelAnnotator(BaseAnnotator): + """ + A class for putting text on an image using provided detections. + + Attributes: + color (Union[Color, ColorPalette]): The color to text on the image, can be a single color or a color palette + """ def __init__(self, color: Union[Color, ColorPalette] = ColorPalette.default(), thickness: int = 2, @@ -173,6 +182,7 @@ def annotate(self, scene: np.ndarray, 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` are not provided, corresponding `class_id` will be used as label. skip_label (bool): Is set to `True`, skips bounding box label annotation. + color_by_track (bool): If set then color will be chosen by tracker id if provided Returns: np.ndarray: The image with the bounding boxes drawn on it From d9f6c5c73c69b03b64502ef19606ce9843139c51 Mon Sep 17 00:00:00 2001 From: hd Date: Tue, 4 Jul 2023 20:42:33 +0200 Subject: [PATCH 04/17] tracked detection method added along with trackstorage --- supervision/__init__.py | 18 +- supervision/annotators/composable.py | 98 +++++++-- supervision/annotators/core.py | 298 +++++++++++++++++++++------ supervision/detection/track.py | 70 +++++++ 4 files changed, 405 insertions(+), 79 deletions(-) create mode 100644 supervision/detection/track.py diff --git a/supervision/__init__.py b/supervision/__init__.py index 5ccabe3e6..748fc4a15 100644 --- a/supervision/__init__.py +++ b/supervision/__init__.py @@ -1,17 +1,29 @@ __version__ = "0.11.1" +from supervision.annotators.composable import ( + DetectionAnnotator, + SegmentationAnnotator, + TrackAnnotator, + TrackedDetectionAnnotator, +) +from supervision.annotators.core import ( + BoxAnnotator, + BoxMaskAnnotator, + LabelAnnotator, + MaskAnnotator, + PillowLabelAnnotator, + TrackAnnotator, +) from supervision.classification.core import Classifications from supervision.dataset.core import ( BaseDataset, ClassificationDataset, DetectionDataset, ) - -from supervision.annotators.core import LabelAnnotator, MaskAnnotator, BoxAnnotator, PillowLabelAnnotator -from supervision.annotators.composable import DetectionAnnotator, TrackAnnotator, SegmentationAnnotator, TrackedDetectionAnnotator from supervision.detection.core import Detections from supervision.detection.line_counter import LineZone, LineZoneAnnotator from supervision.detection.tools.polygon_zone import PolygonZone, PolygonZoneAnnotator +from supervision.detection.track import TrackStorage from supervision.detection.utils import ( box_iou_batch, filter_polygons_by_area, diff --git a/supervision/annotators/composable.py b/supervision/annotators/composable.py index 58b6a2aae..753f6852f 100644 --- a/supervision/annotators/composable.py +++ b/supervision/annotators/composable.py @@ -3,25 +3,67 @@ import numpy as np -from supervision.annotators.core import BoxAnnotator, LabelAnnotator, MaskAnnotator, TrackAnnotator +from supervision.annotators.core import ( + BoxAnnotator, + LabelAnnotator, + MaskAnnotator, + TrackAnnotator, +) from supervision.detection.core import Detections +from supervision.detection.track import TrackStorage class ComposableAnnotator(ABC): - def __init__(self, ): + def __init__( + self, + ): self.annotators = [] - def annotate(self, scene: np.ndarray, detections: Detections, labels: Optional[List[str]] = None, - color_by_track: bool = False,) -> np.ndarray: + def annotate( + self, + scene: np.ndarray, + detections: Detections, + labels: Optional[List[str]] = None, + color_by_track: bool = False, + ) -> np.ndarray: for annotator in self.annotators: if isinstance(annotator, LabelAnnotator): - scene = annotator.annotate(scene=scene, detections=detections, labels=labels, color_by_track=color_by_track) + scene = annotator.annotate( + scene=scene, + detections=detections, + labels=labels, + color_by_track=color_by_track, + ) + elif isinstance(annotator, TrackAnnotator): + scene = annotator.annotate( + scene=scene, color_by_track=color_by_track + ) else: - scene = annotator.annotate(scene=scene, detections=detections, color_by_track=color_by_track) + scene = annotator.annotate( + scene=scene, detections=detections, color_by_track=color_by_track + ) return scene class DetectionAnnotator(ComposableAnnotator): + """ + Highlevel API for drawing Object Detection output. This will use Box and Label Annotators + Example: + ```python + >>> import supervision as sv + + >>> classes = ['person', ...] + >>> image = ... + >>> detections = sv.Detections(...) + + >>> detection_annotator = sv.DetectionAnnotator() + >>> annotated_frame = detection_annotator.annotate( + ... scene=image.copy(), + ... detections=detections + ... ) + ``` + """ + def __init__(self): super().__init__() self.annotators = [ @@ -32,8 +74,23 @@ def __init__(self): class SegmentationAnnotator(ComposableAnnotator): """ - A class for drawing segmentation mask, bounding box and labels on an image using provided detections. + High level API for drawing segmentation mask, bounding box and labels on an image using provided detections. + Example: + ```python + >>> import supervision as sv + + >>> classes = ['person', ...] + >>> image = ... + >>> detections = sv.Detections(...) + + >>> segmentation_annotator = sv.SegmentationAnnotator() + >>> annotated_frame = segmentation_annotator.annotate( + ... scene=image.copy(), + ... detections=detections + ... ) + ``` """ + def __init__(self): super().__init__() self.annotators = [ @@ -41,17 +98,32 @@ def __init__(self): LabelAnnotator(), MaskAnnotator(), ] - raise NotImplementedError class TrackedDetectionAnnotator(ComposableAnnotator): """ - A class for drawing object trajectories, bounding box and tracker ids on an image using provided detections. - """ - def __init__(self): + A class for drawing object trajectories, bounding box and tracker ids on an image using provided detections. + User have freedom to choose color based on class_ids or tracker_ids + Example: + ```python + >>> import supervision as sv + + >>> classes = ['person', ...] + >>> image = ... + >>> detections = sv.Detections(...) + + >>> tracked_detection_annotator = sv.TrackedDetectionAnnotator() + >>> annotated_frame = tracked_detection_annotator.annotate( + ... scene=image.copy(), + ... detections=detections + ... ) + ``` + """ + + def __init__(self, tracks: TrackStorage): super().__init__() self.annotators = [ BoxAnnotator(), LabelAnnotator(), - TrackAnnotator(), - ] \ No newline at end of file + TrackAnnotator(tracks=tracks), + ] diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index 47a14aeed..58198c099 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -1,10 +1,11 @@ from abc import ABC, abstractmethod -from typing import Union, Optional, Dict, List +from typing import List, Optional, Union -import numpy as np import cv2 +import numpy as np from supervision.detection.core import Detections +from supervision.detection.track import TrackStorage from supervision.draw.color import Color, ColorPalette @@ -18,19 +19,20 @@ class BoxAnnotator(BaseAnnotator): """ Basic bounding box annotation class """ + def __init__( - self, - color: Union[Color, ColorPalette] = ColorPalette.default(), - thickness: int = 2, + self, + color: Union[Color, ColorPalette] = ColorPalette.default(), + thickness: int = 2, ): self.color: Union[Color, ColorPalette] = color self.thickness: int = thickness def annotate( - self, - scene: np.ndarray, - detections: Detections, - color_by_track: bool = False, + self, + scene: np.ndarray, + detections: Detections, + color_by_track: bool = False, ) -> np.ndarray: """ Draws bounding boxes on the frame using the detections provided. @@ -62,7 +64,9 @@ def annotate( if color_by_track: tracker_id = ( - detections.tracker_id[i] if detections.tracker_id is not None else None + detections.tracker_id[i] + if detections.tracker_id is not None + else None ) idx = tracker_id if tracker_id is not None else i else: @@ -103,7 +107,12 @@ def __init__( self.color: Union[Color, ColorPalette] = color self.opacity = opacity - def annotate(self, scene: np.ndarray, detections: Detections, color_by_track: bool = False,) -> np.ndarray: + def annotate( + self, + scene: np.ndarray, + detections: Detections, + color_by_track: bool = False, + ) -> np.ndarray: """ Overlays the masks on the given image based on the provided detections, with a specified opacity. @@ -113,6 +122,20 @@ def annotate(self, scene: np.ndarray, detections: Detections, color_by_track: bo Returns: np.ndarray: The image with the masks overlaid + Example: + ```python + >>> import supervision as sv + + >>> classes = ['person', ...] + >>> image = ... + >>> detections = sv.Detections(...) + + >>> mask_annotator = sv.MaskAnnotator() + >>> annotated_frame = mask_annotator.annotate( + ... scene=image.copy(), + ... detections=detections + ... ) + ``` """ if detections.mask is None: return scene @@ -120,7 +143,9 @@ def annotate(self, scene: np.ndarray, detections: Detections, color_by_track: bo for i in np.flip(np.argsort(detections.area)): if color_by_track: tracker_id = ( - detections.tracker_id[i] if detections.tracker_id is not None else None + detections.tracker_id[i] + if detections.tracker_id is not None + else None ) idx = tracker_id if tracker_id is not None else i else: @@ -156,12 +181,15 @@ class LabelAnnotator(BaseAnnotator): color (Union[Color, ColorPalette]): The color to text on the image, can be a single color or a color palette """ - def __init__(self, color: Union[Color, ColorPalette] = ColorPalette.default(), + 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,): + text_padding: int = 10, + ): self.color: Union[Color, ColorPalette] = color self.thickness: int = thickness self.text_color: Color = text_color @@ -169,62 +197,69 @@ def __init__(self, color: Union[Color, ColorPalette] = ColorPalette.default(), self.text_thickness: int = text_thickness self.text_padding: int = text_padding - def annotate(self, scene: np.ndarray, + def annotate( + self, + scene: np.ndarray, detections: Detections, labels: Optional[List[str]] = None, - skip_label: bool = False, - color_by_track: bool = False,) -> np.ndarray: + color_by_track: bool = False, + ) -> np.ndarray: + """ + Draws text on the frame using the detections provided and label. + + Args: + 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` are not provided, corresponding `class_id` will be used as label. + color_by_track (bool): If set then color will be chosen by tracker id if provided + Returns: + np.ndarray: The image with the bounding boxes drawn on it + + Example: + ```python + >>> import supervision as sv + + >>> classes = ['person', ...] + >>> image = ... + >>> detections = sv.Detections(...) + + >>> label_annotator = sv.LabelAnnotator() + >>> labels = [ + ... f"{classes[class_id]} {confidence:0.2f}" + ... for _, _, confidence, class_id, _ + ... in detections + ... ] + >>> annotated_frame = label_annotator.annotate( + ... scene=image.copy(), + ... detections=detections, + ... labels=labels + ... ) + ``` """ - Draws bounding boxes on the frame using the detections provided. - - Args: - 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` are not provided, corresponding `class_id` will be used as label. - skip_label (bool): Is set to `True`, skips bounding box label annotation. - color_by_track (bool): If set then color will be chosen by tracker id if provided - Returns: - np.ndarray: The image with the bounding boxes drawn on it - - Example: - ```python - >>> import supervision as sv - - >>> classes = ['person', ...] - >>> image = ... - >>> detections = sv.Detections(...) - - >>> label_annotator = sv.LabelAnnotator() - >>> labels = [ - ... f"{classes[class_id]} {confidence:0.2f}" - ... for _, _, confidence, class_id, _ - ... in detections - ... ] - >>> annotated_frame = label_annotator.annotate( - ... scene=image.copy(), - ... detections=detections, - ... labels=labels - ... ) - ``` - """ font = cv2.FONT_HERSHEY_SIMPLEX for i in range(len(detections)): x1, y1, x2, y2 = detections.xyxy[i].astype(int) - class_id = ( - detections.class_id[i] if detections.class_id is not None else None - ) - idx = class_id if class_id is not None else i + if color_by_track: + tracker_id = ( + detections.tracker_id[i] + if detections.tracker_id is not None + else None + ) + idx = tracker_id if tracker_id is not None else i + else: + class_id = ( + detections.class_id[i] if detections.class_id is not None else None + ) + idx = class_id if class_id is not None else i + color = ( self.color.by_idx(idx) if isinstance(self.color, ColorPalette) else self.color ) - if skip_label: - continue - text = ( - f"{class_id}" + f"{idx}" if (labels is None or len(detections) != len(labels)) else labels[i] ) @@ -265,16 +300,153 @@ def annotate(self, scene: np.ndarray, return scene -class PillowLabelAnnotator(BaseAnnotator): - def annotate(self, scene: np.ndarray, detections: Detections) -> np.ndarray: - pass +class BoxMaskAnnotator(BaseAnnotator): + def __init__( + self, + color: Union[Color, ColorPalette] = ColorPalette.default(), + opacity: float = 0.5, + ): + self.color: Union[Color, ColorPalette] = color + self.opacity = opacity + + def annotate( + self, + scene: np.ndarray, + detections: Detections, + color_by_track: bool = False, + ) -> np.ndarray: + """ + Overlays the rectangle masks on the given image based on the provided detections, with a specified opacity. + + Args: + scene (np.ndarray): The image on which the masks will be overlaid + detections (Detections): The detections for which the masks will be overlaid + color_by_track (bool): If set then color will be chosen by tracker id if provided + Returns: + np.ndarray: The image with the masks overlaid + Example: + ```python + >>> import supervision as sv + + >>> classes = ['person', ...] + >>> image = ... + >>> detections = sv.Detections(...) + + >>> box_mask_annotator = sv.MaskAnnotator() + >>> annotated_frame = box_mask_annotator.annotate( + ... scene=image.copy(), + ... detections=detections + ... ) + ``` + """ + overlay_img = np.zeros_like(scene, np.uint8) + for i in range(len(detections)): + x1, y1, x2, y2 = detections.xyxy[i].astype(int) + if color_by_track: + tracker_id = ( + detections.tracker_id[i] + if detections.tracker_id is not None + else None + ) + idx = tracker_id if tracker_id is not None else i + else: + class_id = ( + detections.class_id[i] if detections.class_id is not None else None + ) + idx = class_id if class_id is not None else i + color = ( + self.color.by_idx(idx) + if isinstance(self.color, ColorPalette) + else self.color + ) + cv2.rectangle( + img=overlay_img, + pt1=(x1, y1), + pt2=(x2, y2), + color=color.as_bgr(), + thickness=-1, + ) + + mask = overlay_img.astype(bool) + scene[mask] = cv2.addWeighted( + scene, self.opacity, overlay_img, 1 - self.opacity, 0 + )[mask] + return scene class TrackAnnotator(BaseAnnotator): - def annotate(self, scene: np.ndarray, detections: Detections) -> np.ndarray: - pass + """ + Initialize TrackAnnotator + """ + + def __init__( + self, + tracks: TrackStorage, + color: Union[Color, ColorPalette] = ColorPalette.default(), + thickness: int = 2, + ): + self.track_storage = tracks + self.color: Union[Color, ColorPalette] = color + self.thickness: int = thickness + self.boundry_tolerance = 20 + def annotate(self, scene: np.ndarray, color_by_track: bool = False) -> np.ndarray: + """ + Draws the object trajectory on the frame using the trace provided. + Attributes: + scene (np.ndarray): The image on which the object trajectories will be drawn. + trace (Trace): The trace that will be used to draw the previous and current position. -class BoxMaskAnnotator(BaseAnnotator): + Returns: + np.ndarray: The image with the object trajectories on it. + ```python + >>> import supervision as sv + >>> track_storage = sv.TrackStorage() + >>> track_annotator = sv.TrackAnnotator(track_storage) + >>> for frame in sv.get_video_frames_generator(source_path='source_video.mp4'): + >>> detections = sv.Detections(...) + >>> tracked_objects = tracker(...) + >>> tracked_detections = sv.Detections(tracked_objects) + >>> track_storage.update(tracked_detections) + >>> track_annotator.annotate(scene) + """ + img_h, img_w, _ = scene.shape + unique_ids = np.unique(self.track_storage.storage[:, -1]) + for unique_id in unique_ids: + valid = np.where(self.track_storage.storage[:, -1] == unique_id)[0] + + frames = self.track_storage.storage[valid, 0] + latest_frame = np.argmax(frames) + points_to_draw = self.track_storage.storage[valid, 1:3] + + n_pts = points_to_draw.shape[0] + headx, heady = int(points_to_draw[latest_frame][0]), int( + points_to_draw[latest_frame][1] + ) + + if headx > self.boundry_tolerance and heady > self.boundry_tolerance: + if color_by_track: + idx = int(unique_id) + else: + idx = int(self.track_storage.storage[0, -2]) + color = ( + self.color.by_idx(idx) + if isinstance(self.color, ColorPalette) + else self.color + ) + + for i in range(n_pts - 1): + px, py = int(points_to_draw[i][0]), int(points_to_draw[i][1]) + qx, qy = int(points_to_draw[i + 1][0]), int( + points_to_draw[i + 1][1] + ) + cv2.line(scene, (px, py), (qx, qy), color.as_bgr(), self.thickness) + scene = cv2.circle( + scene, (headx, heady), int(10), color.as_bgr(), thickness=-1 + ) + return scene + + +class PillowLabelAnnotator(BaseAnnotator): def annotate(self, scene: np.ndarray, detections: Detections) -> np.ndarray: pass diff --git a/supervision/detection/track.py b/supervision/detection/track.py new file mode 100644 index 000000000..8b40e38a2 --- /dev/null +++ b/supervision/detection/track.py @@ -0,0 +1,70 @@ +import numpy as np + +from supervision.detection.core import Detections +from supervision.geometry.core import Position + + +class TrackStorage: + """ + Trace the object trajectory with Detections with tracker id + """ + + def __init__( + self, + position: Position = Position.CENTER, + max_length: int = 10, + ): + """ + Initialize the track storage to store tracked detections from tracker + Args: + position [sv.Position]: Trace position whethere mid-point or bottom center point + max_length [int]: Length of previous detections to annotate + Detections objects are stored as numpy ndarray with #counter, x, y, class, track_id + + Example: + ```python + >>> import supervision as sv + >>> track_storage = sv.TrackStorage() + >>> for frame in sv.get_video_frames_generator(source_path='source_video.mp4'): + >>> detections = sv.Detections(...) + >>> tracked_objects = tracker(...) + >>> tracked_detections = sv.Detections(tracked_objects) + >>> track_storage.update(tracked_detections) + """ + self.position = position + self.max_length = max_length + self.frame_counter = 0 + self.storage = np.zeros((0, 5)) + + def update(self, frame_counter: int, detections: Detections) -> None: + if detections.xyxy.shape[0] > 0: + xyxy = detections.xyxy + + if self.position == Position.CENTER: + x = (xyxy[:, 0] + xyxy[:, 2]) / 2 + y = (xyxy[:, 1] + xyxy[:, 3]) / 2 + elif self.position == Position.BOTTOM_CENTER: + x = (xyxy[:, 0] + xyxy[:, 2]) / 2 + y = xyxy[:, 3] + + new_detections = np.zeros(shape=(xyxy.shape[0], 5)) + new_detections[:, 0] = frame_counter + new_detections[:, 1] = x + new_detections[:, 2] = y + new_detections[:, 3] = detections.class_id + new_detections[:, 4] = detections.tracker_id + self.storage = np.append(self.storage, new_detections, axis=0) + self.frame_counter = frame_counter + self._remove_lost(tracker_ids=detections.tracker_id) + self._remove_previous() + + def _remove_previous(self) -> None: + to_remove_frames = self.frame_counter - self.max_length + if self.storage.shape[0] > 0: + self.storage = self.storage[ + self.storage[:, 0] > to_remove_frames + ] + + def _remove_lost(self, tracker_ids) -> None: + if self.storage.shape[0] > 0: + self.storage = self.storage[np.isin(self.storage[:, -1], tracker_ids)] From d5b8a3aa1a85c12ab0413083f76b683cf48fe104 Mon Sep 17 00:00:00 2001 From: hd Date: Wed, 5 Jul 2023 15:01:16 +0200 Subject: [PATCH 05/17] Adding cornered box annotator and pillow label annotator --- setup.py | 3 +- supervision/__init__.py | 1 + supervision/annotators/composable.py | 4 +- supervision/annotators/core.py | 245 ++++++++++++++++++++++++++- supervision/detection/track.py | 4 +- 5 files changed, 248 insertions(+), 9 deletions(-) diff --git a/setup.py b/setup.py index ac3add4fb..69ca93b6e 100644 --- a/setup.py +++ b/setup.py @@ -27,7 +27,8 @@ def get_version(): 'numpy>=1.20.0', 'opencv-python', 'matplotlib', - 'pyyaml' + 'pyyaml', + 'pillow' ], packages=find_packages(exclude=("tests",)), extras_require={ diff --git a/supervision/__init__.py b/supervision/__init__.py index 748fc4a15..797a5596e 100644 --- a/supervision/__init__.py +++ b/supervision/__init__.py @@ -9,6 +9,7 @@ from supervision.annotators.core import ( BoxAnnotator, BoxMaskAnnotator, + CorneredBoxAnotator, LabelAnnotator, MaskAnnotator, PillowLabelAnnotator, diff --git a/supervision/annotators/composable.py b/supervision/annotators/composable.py index 753f6852f..a0fd6718e 100644 --- a/supervision/annotators/composable.py +++ b/supervision/annotators/composable.py @@ -35,9 +35,7 @@ def annotate( color_by_track=color_by_track, ) elif isinstance(annotator, TrackAnnotator): - scene = annotator.annotate( - scene=scene, color_by_track=color_by_track - ) + scene = annotator.annotate(scene=scene, color_by_track=color_by_track) else: scene = annotator.annotate( scene=scene, detections=detections, color_by_track=color_by_track diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index 58198c099..0bc684f9a 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -1,8 +1,11 @@ +import os.path from abc import ABC, abstractmethod from typing import List, Optional, Union import cv2 import numpy as np +import PIL.Image +from PIL import Image, ImageDraw, ImageFont from supervision.detection.core import Detections from supervision.detection.track import TrackStorage @@ -448,5 +451,243 @@ def annotate(self, scene: np.ndarray, color_by_track: bool = False) -> np.ndarra class PillowLabelAnnotator(BaseAnnotator): - def annotate(self, scene: np.ndarray, detections: Detections) -> np.ndarray: - pass + def __init__( + self, + color: Union[Color, ColorPalette] = ColorPalette.default(), + text_color: Color = Color.black(), + text_padding: int = 20, + ): + self.font = ImageFont.load_default() + self.color: Union[Color, ColorPalette] = color + self.text_color: Color = text_color + self.text_padding: int = text_padding + + def annotate( + self, + scene: np.ndarray, + detections: Detections, + labels: Optional[List[str]] = None, + color_by_track: bool = False, + font: Optional[str] = None, + font_size: Optional[int] = 15, + ) -> np.ndarray: + """ + Draws text on the frame using the detections provided and label. + + Args: + 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` are not provided, corresponding `class_id` will be used as label. + color_by_track (bool): If set then color will be chosen by tracker id if provided + Returns: + np.ndarray: The image with the bounding boxes drawn on it + + Example: + ```python + >>> import supervision as sv + + >>> classes = ['person', ...] + >>> image = ... + >>> detections = sv.Detections(...) + + >>> pil_label_annotator = sv.PillowLabelAnnotator() + >>> labels = [ + ... f"{classes[class_id]} {confidence:0.2f}" + ... for _, _, confidence, class_id, _ + ... in detections + ... ] + >>> annotated_frame = pil_label_annotator.annotate( + ... scene=image.copy(), + ... detections=detections, + ... labels=labels, + ... font=FONT_FILE_PATH, + ... ) + ``` + """ + if font and os.path.exists(font): + self.font = ImageFont.truetype(font, font_size) + + pil_image = Image.fromarray(scene) + draw = ImageDraw.Draw(pil_image) + text_color = "#fff" + + for i in range(len(detections)): + x1, y1, x2, y2 = detections.xyxy[i].astype(int) + if color_by_track: + tracker_id = ( + detections.tracker_id[i] + if detections.tracker_id is not None + else None + ) + idx = tracker_id if tracker_id is not None else i + else: + class_id = ( + detections.class_id[i] if detections.class_id is not None else None + ) + idx = class_id if class_id is not None else i + + color = ( + self.color.by_idx(idx) + if isinstance(self.color, ColorPalette) + else self.color + ) + + text = ( + f"{idx}" + if (labels is None or len(detections) != len(labels)) + else labels[i] + ) + + text_bbox = draw.textbbox((x1, y1), text, font=self.font) + + text_height = text_bbox[3] - text_bbox[1] + text_width = text_bbox[2] - text_bbox[0] + + text_x = x1 + self.text_padding / 2 + text_y = y1 - self.text_padding / 2 - text_height + + text_background_x1 = x1 + text_background_y1 = y1 - self.text_padding / 2 - text_height + + text_background_x2 = x1 + 2 * self.text_padding / 2 + text_width + text_background_y2 = y1 # correct + + draw.rectangle( + ( + text_background_x1, + text_background_y1, + text_background_x2, + text_background_y2, + ), + fill=color.as_bgr(), + ) + draw.text((text_x, text_y), text, font=self.font, fill=text_color) + + scene = np.asarray(pil_image) + return scene + + +class CorneredBoxAnotator(BaseAnnotator): + def __init__( + self, + color: Union[Color, ColorPalette] = ColorPalette.default(), + thickness: int = 2, + ): + self.color: Union[Color, ColorPalette] = color + self.thickness: int = thickness + + def annotate( + self, + scene: np.ndarray, + detections: Detections, + color_by_track: bool = False, + ): + """ + Draws cornered bounding boxes on the frame using the detections provided. + Args: + 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` are not provided, corresponding `class_id` will be used as 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 + Example: + ```python + >>> import supervision as sv + >>> classes = ['person', ...] + >>> image = ... + >>> detections = sv.Detections(...) + >>> corner_box_annotator = sv.CorneredBoxAnotator() + >>> annotated_frame = corner_box_annotator.annotate( + ... scene=image.copy(), + ... detections=detections, + ... labels=labels + ... ) + ``` + """ + + line_thickness = self.thickness + 2 + for i in range(len(detections)): + x1, y1, x2, y2 = detections.xyxy[i].astype(int) + if color_by_track: + tracker_id = ( + detections.tracker_id[i] + if detections.tracker_id is not None + else None + ) + idx = tracker_id if tracker_id is not None else i + else: + class_id = ( + detections.class_id[i] if detections.class_id is not None else None + ) + idx = class_id if class_id is not None else i + color = ( + self.color.by_idx(idx) + if isinstance(self.color, ColorPalette) + else self.color + ) + + box_width = x2 - x1 + box_height = y2 - y1 + + cv2.line( + scene, + (x1, y1), + (x1 + int(0.2 * box_width), y1), + color.as_bgr(), + thickness=line_thickness, + ) + cv2.line( + scene, + (x2 - int(0.2 * box_width), y1), + (x2, y1), + color.as_bgr(), + thickness=line_thickness, + ) + + cv2.line( + scene, + (x1, y2), + (x1 + int(0.2 * box_width), y2), + color.as_bgr(), + thickness=line_thickness, + ) + cv2.line( + scene, + (x2 - int(0.2 * box_width), y2), + (x2, y2), + color.as_bgr(), + thickness=line_thickness, + ) + # + cv2.line( + scene, + (x1, y1), + (x1, y1 + int(0.2 * box_height)), + color.as_bgr(), + thickness=line_thickness, + ) + cv2.line( + scene, + (x2, y1 + int(0.2 * box_height)), + (x2, y1), + color.as_bgr(), + thickness=line_thickness, + ) + + cv2.line( + scene, + (x1, y2 - int(0.2 * box_height)), + (x1, y2), + color.as_bgr(), + thickness=line_thickness, + ) + cv2.line( + scene, + (x2, y2 - int(0.2 * box_height)), + (x2, y2), + color.as_bgr(), + thickness=line_thickness, + ) + + return scene diff --git a/supervision/detection/track.py b/supervision/detection/track.py index 8b40e38a2..c0d5a6872 100644 --- a/supervision/detection/track.py +++ b/supervision/detection/track.py @@ -61,9 +61,7 @@ def update(self, frame_counter: int, detections: Detections) -> None: def _remove_previous(self) -> None: to_remove_frames = self.frame_counter - self.max_length if self.storage.shape[0] > 0: - self.storage = self.storage[ - self.storage[:, 0] > to_remove_frames - ] + self.storage = self.storage[self.storage[:, 0] > to_remove_frames] def _remove_lost(self, tracker_ids) -> None: if self.storage.shape[0] > 0: From c705d50e742d542e84d28b9dcec70e8679ae38d6 Mon Sep 17 00:00:00 2001 From: hd Date: Wed, 5 Jul 2023 15:39:22 +0200 Subject: [PATCH 06/17] Marked as ready to review --- supervision/annotators/core.py | 5 +- supervision/detection/annotate.py | 173 ------------------------------ 2 files changed, 2 insertions(+), 176 deletions(-) diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index 0bc684f9a..e1b48bdbb 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -240,6 +240,7 @@ def annotate( ``` """ font = cv2.FONT_HERSHEY_SIMPLEX + for i in range(len(detections)): x1, y1, x2, y2 = detections.xyxy[i].astype(int) if color_by_track: @@ -398,7 +399,7 @@ def annotate(self, scene: np.ndarray, color_by_track: bool = False) -> np.ndarra Draws the object trajectory on the frame using the trace provided. Attributes: scene (np.ndarray): The image on which the object trajectories will be drawn. - trace (Trace): The trace that will be used to draw the previous and current position. + tracks (TrackStorage): The track storage that will be used to draw the previous and current position. Returns: np.ndarray: The image with the object trajectories on it. @@ -587,8 +588,6 @@ def annotate( Args: 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` are not provided, corresponding `class_id` will be used as 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 Example: diff --git a/supervision/detection/annotate.py b/supervision/detection/annotate.py index 17c78df32..3b3e40da6 100644 --- a/supervision/detection/annotate.py +++ b/supervision/detection/annotate.py @@ -37,179 +37,6 @@ def __init__( self.text_thickness: int = text_thickness self.text_padding: int = text_padding - def annotate_advance( - self, - scene: np.ndarray, - detections: Detections, - labels: Optional[List[str]] = None, - skip_label: bool = False, - ): - """ - Draws bounding boxes on the frame using the detections provided. - - Args: - 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` are not provided, corresponding `class_id` will be used as 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 - - Example: - ```python - >>> import supervision as sv - - >>> classes = ['person', ...] - >>> image = ... - >>> detections = sv.Detections(...) - - >>> box_annotator = sv.BoxAnnotator() - >>> labels = [ - ... f"{classes[class_id]} {confidence:0.2f}" - ... for _, _, confidence, class_id, _ - ... in detections - ... ] - >>> annotated_frame = box_annotator.annotate( - ... scene=image.copy(), - ... detections=detections, - ... labels=labels - ... ) - ``` - """ - font = cv2.FONT_HERSHEY_SIMPLEX - overlay_img = np.zeros_like(scene, np.uint8) - alpha = 0.8 - - line_thickness = self.thickness + 2 - for i in range(len(detections)): - x1, y1, x2, y2 = detections.xyxy[i].astype(int) - class_id = ( - detections.class_id[i] if detections.class_id is not None else None - ) - idx = class_id if class_id is not None else i - color = ( - self.color.by_idx(idx) - if isinstance(self.color, ColorPalette) - else self.color - ) - - box_width = x2 - x1 - box_height = y2 - y1 - - cv2.line( - scene, - (x1, y1), - (x1 + int(0.2 * box_width), y1), - color.as_bgr(), - thickness=line_thickness, - ) - cv2.line( - scene, - (x2 - int(0.2 * box_width), y1), - (x2, y1), - color.as_bgr(), - thickness=line_thickness, - ) - - cv2.line( - scene, - (x1, y2), - (x1 + int(0.2 * box_width), y2), - color.as_bgr(), - thickness=line_thickness, - ) - cv2.line( - scene, - (x2 - int(0.2 * box_width), y2), - (x2, y2), - color.as_bgr(), - thickness=line_thickness, - ) - # - cv2.line( - scene, - (x1, y1), - (x1, y1 + int(0.2 * box_height)), - color.as_bgr(), - thickness=line_thickness, - ) - cv2.line( - scene, - (x2, y1 + int(0.2 * box_height)), - (x2, y1), - color.as_bgr(), - thickness=line_thickness, - ) - - cv2.line( - scene, - (x1, y2 - int(0.2 * box_height)), - (x1, y2), - color.as_bgr(), - thickness=line_thickness, - ) - cv2.line( - scene, - (x2, y2 - int(0.2 * box_height)), - (x2, y2), - color.as_bgr(), - thickness=line_thickness, - ) - - cv2.rectangle( - img=overlay_img, - pt1=(x1, y1), - pt2=(x2, y2), - color=color.as_bgr(), - thickness=-1, - ) - 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, - ) - - mask = overlay_img.astype(bool) - scene[mask] = cv2.addWeighted(scene, alpha, overlay_img, 1 - alpha, 0)[mask] - return scene - def annotate( self, scene: np.ndarray, From 2a2c259964e00923dd57b8ff545b1317cae889fe Mon Sep 17 00:00:00 2001 From: hd Date: Fri, 7 Jul 2023 09:37:03 +0200 Subject: [PATCH 07/17] :fire: Consistent code style --- supervision/annotators/composable.py | 48 +++++++++++++--------- supervision/annotators/core.py | 59 +++++++++++++++------------- 2 files changed, 61 insertions(+), 46 deletions(-) diff --git a/supervision/annotators/composable.py b/supervision/annotators/composable.py index a0fd6718e..a3db9e7c1 100644 --- a/supervision/annotators/composable.py +++ b/supervision/annotators/composable.py @@ -24,7 +24,6 @@ def annotate( scene: np.ndarray, detections: Detections, labels: Optional[List[str]] = None, - color_by_track: bool = False, ) -> np.ndarray: for annotator in self.annotators: if isinstance(annotator, LabelAnnotator): @@ -32,14 +31,11 @@ def annotate( scene=scene, detections=detections, labels=labels, - color_by_track=color_by_track, ) elif isinstance(annotator, TrackAnnotator): - scene = annotator.annotate(scene=scene, color_by_track=color_by_track) + scene = annotator.annotate(scene=scene) else: - scene = annotator.annotate( - scene=scene, detections=detections, color_by_track=color_by_track - ) + scene = annotator.annotate(scene=scene, detections=detections) return scene @@ -62,12 +58,15 @@ class DetectionAnnotator(ComposableAnnotator): ``` """ - def __init__(self): + def __init__( + self, + color_by_track: bool = False, + skip_label: bool = False, + ): super().__init__() - self.annotators = [ - BoxAnnotator(), - LabelAnnotator(), - ] + self.annotators = [BoxAnnotator(color_by_track=color_by_track)] + if not skip_label: + self.annotators.append(LabelAnnotator(color_by_track=color_by_track)) class SegmentationAnnotator(ComposableAnnotator): @@ -89,13 +88,18 @@ class SegmentationAnnotator(ComposableAnnotator): ``` """ - def __init__(self): + def __init__( + self, + color_by_track: bool = False, + skip_label: bool = False, + ): super().__init__() self.annotators = [ - BoxAnnotator(), - LabelAnnotator(), - MaskAnnotator(), + BoxAnnotator(color_by_track=color_by_track), + MaskAnnotator(color_by_track=color_by_track), ] + if not skip_label: + self.annotators.append(LabelAnnotator(color_by_track=color_by_track)) class TrackedDetectionAnnotator(ComposableAnnotator): @@ -118,10 +122,16 @@ class TrackedDetectionAnnotator(ComposableAnnotator): ``` """ - def __init__(self, tracks: TrackStorage): + def __init__( + self, + tracks: TrackStorage, + color_by_track: bool = False, + skip_label: bool = False, + ): super().__init__() self.annotators = [ - BoxAnnotator(), - LabelAnnotator(), - TrackAnnotator(tracks=tracks), + BoxAnnotator(color_by_track=color_by_track), + TrackAnnotator(tracks=tracks, color_by_track=color_by_track), ] + if not skip_label: + self.annotators.append(LabelAnnotator(color_by_track=color_by_track)) diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index e1b48bdbb..a6932926f 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -4,7 +4,6 @@ import cv2 import numpy as np -import PIL.Image from PIL import Image, ImageDraw, ImageFont from supervision.detection.core import Detections @@ -27,15 +26,16 @@ def __init__( self, color: Union[Color, ColorPalette] = ColorPalette.default(), thickness: int = 2, + color_by_track: bool = False, ): self.color: Union[Color, ColorPalette] = color self.thickness: int = thickness + self.color_by_track = color_by_track def annotate( self, scene: np.ndarray, detections: Detections, - color_by_track: bool = False, ) -> np.ndarray: """ Draws bounding boxes on the frame using the detections provided. @@ -43,7 +43,6 @@ def annotate( Args: 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 - color_by_track (bool): It allows to pick color by tracker id if present Returns: np.ndarray: The image with the bounding boxes drawn on it @@ -65,7 +64,7 @@ def annotate( for i in range(len(detections)): x1, y1, x2, y2 = detections.xyxy[i].astype(int) - if color_by_track: + if self.color_by_track: tracker_id = ( detections.tracker_id[i] if detections.tracker_id is not None @@ -106,15 +105,16 @@ def __init__( self, color: Union[Color, ColorPalette] = ColorPalette.default(), opacity: float = 0.5, + color_by_track: bool = False, ): self.color: Union[Color, ColorPalette] = color self.opacity = opacity + self.color_by_track = color_by_track def annotate( self, scene: np.ndarray, detections: Detections, - color_by_track: bool = False, ) -> np.ndarray: """ Overlays the masks on the given image based on the provided detections, with a specified opacity. @@ -144,7 +144,7 @@ def annotate( return scene for i in np.flip(np.argsort(detections.area)): - if color_by_track: + if self.color_by_track: tracker_id = ( detections.tracker_id[i] if detections.tracker_id is not None @@ -192,6 +192,7 @@ def __init__( text_scale: float = 0.5, text_thickness: int = 1, text_padding: int = 10, + color_by_track: bool = False, ): self.color: Union[Color, ColorPalette] = color self.thickness: int = thickness @@ -199,13 +200,13 @@ def __init__( self.text_scale: float = text_scale self.text_thickness: int = text_thickness self.text_padding: int = text_padding + self.color_by_track = color_by_track def annotate( self, scene: np.ndarray, detections: Detections, labels: Optional[List[str]] = None, - color_by_track: bool = False, ) -> np.ndarray: """ Draws text on the frame using the detections provided and label. @@ -214,7 +215,6 @@ def annotate( 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` are not provided, corresponding `class_id` will be used as label. - color_by_track (bool): If set then color will be chosen by tracker id if provided Returns: np.ndarray: The image with the bounding boxes drawn on it @@ -243,7 +243,7 @@ def annotate( for i in range(len(detections)): x1, y1, x2, y2 = detections.xyxy[i].astype(int) - if color_by_track: + if self.color_by_track: tracker_id = ( detections.tracker_id[i] if detections.tracker_id is not None @@ -309,15 +309,16 @@ def __init__( self, color: Union[Color, ColorPalette] = ColorPalette.default(), opacity: float = 0.5, + color_by_track: bool = False, ): self.color: Union[Color, ColorPalette] = color self.opacity = opacity + self.color_by_track = color_by_track def annotate( self, scene: np.ndarray, detections: Detections, - color_by_track: bool = False, ) -> np.ndarray: """ Overlays the rectangle masks on the given image based on the provided detections, with a specified opacity. @@ -325,7 +326,6 @@ def annotate( Args: scene (np.ndarray): The image on which the masks will be overlaid detections (Detections): The detections for which the masks will be overlaid - color_by_track (bool): If set then color will be chosen by tracker id if provided Returns: np.ndarray: The image with the masks overlaid Example: @@ -346,7 +346,7 @@ def annotate( overlay_img = np.zeros_like(scene, np.uint8) for i in range(len(detections)): x1, y1, x2, y2 = detections.xyxy[i].astype(int) - if color_by_track: + if self.color_by_track: tracker_id = ( detections.tracker_id[i] if detections.tracker_id is not None @@ -388,13 +388,18 @@ def __init__( tracks: TrackStorage, color: Union[Color, ColorPalette] = ColorPalette.default(), thickness: int = 2, + color_by_track: bool = False, ): self.track_storage = tracks self.color: Union[Color, ColorPalette] = color self.thickness: int = thickness self.boundry_tolerance = 20 + self.color_by_track = color_by_track - def annotate(self, scene: np.ndarray, color_by_track: bool = False) -> np.ndarray: + def annotate( + self, + scene: np.ndarray, + ) -> np.ndarray: """ Draws the object trajectory on the frame using the trace provided. Attributes: @@ -429,7 +434,7 @@ def annotate(self, scene: np.ndarray, color_by_track: bool = False) -> np.ndarra ) if headx > self.boundry_tolerance and heady > self.boundry_tolerance: - if color_by_track: + if self.color_by_track: idx = int(unique_id) else: idx = int(self.track_storage.storage[0, -2]) @@ -457,20 +462,24 @@ def __init__( color: Union[Color, ColorPalette] = ColorPalette.default(), text_color: Color = Color.black(), text_padding: int = 20, + color_by_track: bool = False, + font: Optional[str] = None, + font_size: Optional[int] = 15, ): - self.font = ImageFont.load_default() + if font and os.path.exists(font): + self.font = ImageFont.truetype(font, font_size) + else: + self.font = ImageFont.load_default() self.color: Union[Color, ColorPalette] = color self.text_color: Color = text_color self.text_padding: int = text_padding + self.color_by_track = color_by_track def annotate( self, scene: np.ndarray, detections: Detections, labels: Optional[List[str]] = None, - color_by_track: bool = False, - font: Optional[str] = None, - font_size: Optional[int] = 15, ) -> np.ndarray: """ Draws text on the frame using the detections provided and label. @@ -479,7 +488,6 @@ def annotate( 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` are not provided, corresponding `class_id` will be used as label. - color_by_track (bool): If set then color will be chosen by tracker id if provided Returns: np.ndarray: The image with the bounding boxes drawn on it @@ -501,20 +509,16 @@ def annotate( ... scene=image.copy(), ... detections=detections, ... labels=labels, - ... font=FONT_FILE_PATH, ... ) ``` """ - if font and os.path.exists(font): - self.font = ImageFont.truetype(font, font_size) - pil_image = Image.fromarray(scene) draw = ImageDraw.Draw(pil_image) text_color = "#fff" for i in range(len(detections)): x1, y1, x2, y2 = detections.xyxy[i].astype(int) - if color_by_track: + if self.color_by_track: tracker_id = ( detections.tracker_id[i] if detections.tracker_id is not None @@ -573,15 +577,16 @@ def __init__( self, color: Union[Color, ColorPalette] = ColorPalette.default(), thickness: int = 2, + color_by_track: bool = False, ): self.color: Union[Color, ColorPalette] = color self.thickness: int = thickness + self.color_by_track = color_by_track def annotate( self, scene: np.ndarray, detections: Detections, - color_by_track: bool = False, ): """ Draws cornered bounding boxes on the frame using the detections provided. @@ -608,7 +613,7 @@ def annotate( line_thickness = self.thickness + 2 for i in range(len(detections)): x1, y1, x2, y2 = detections.xyxy[i].astype(int) - if color_by_track: + if self.color_by_track: tracker_id = ( detections.tracker_id[i] if detections.tracker_id is not None @@ -658,7 +663,7 @@ def annotate( color.as_bgr(), thickness=line_thickness, ) - # + cv2.line( scene, (x1, y1), From 27efbc044dba6f4700f6195e70baf679f9dc3449 Mon Sep 17 00:00:00 2001 From: hd Date: Fri, 7 Jul 2023 19:18:12 +0200 Subject: [PATCH 08/17] Adding EllipseAnotator and added backward compatibility for box and mask annotator --- supervision/__init__.py | 1 + supervision/annotators/core.py | 82 ++++++++++++++++++++++++++++++++++ 2 files changed, 83 insertions(+) diff --git a/supervision/__init__.py b/supervision/__init__.py index 797a5596e..21f3c81a2 100644 --- a/supervision/__init__.py +++ b/supervision/__init__.py @@ -10,6 +10,7 @@ BoxAnnotator, BoxMaskAnnotator, CorneredBoxAnotator, + EllipseAnotator, LabelAnnotator, MaskAnnotator, PillowLabelAnnotator, diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index a6932926f..94be97e61 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -36,6 +36,7 @@ def annotate( self, scene: np.ndarray, detections: Detections, + **kwargs ) -> np.ndarray: """ Draws bounding boxes on the frame using the detections provided. @@ -115,6 +116,7 @@ def annotate( self, scene: np.ndarray, detections: Detections, + **kwargs ) -> np.ndarray: """ Overlays the masks on the given image based on the provided detections, with a specified opacity. @@ -695,3 +697,83 @@ def annotate( ) return scene + + +class EllipseAnotator(BaseAnnotator): + def __init__( + self, + color: Union[Color, ColorPalette] = ColorPalette.default(), + thickness: int = 2, + color_by_track: bool = False, + start_angle: int = -45, + end_angle: int = 360, + ): + self.color: Union[Color, ColorPalette] = color + self.thickness: int = thickness + self.color_by_track = color_by_track + self.start_angle = start_angle + self.end_angle = end_angle + + def annotate( + self, + scene: np.ndarray, + detections: Detections, + ): + """ + Draws ellipse at bottom of the objects on the frame using the detections provided. + Args: + 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 + Returns: + np.ndarray: The image with the bounding boxes drawn on it + Example: + ```python + >>> import supervision as sv + >>> classes = ['person', ...] + >>> image = ... + >>> detections = sv.Detections(...) + >>> ellipse_annotator = sv.EllipseAnotator() + >>> annotated_frame = ellipse_annotator.annotate( + ... scene=image.copy(), + ... detections=detections, + ... labels=labels + ... ) + ``` + """ + + for i in range(len(detections)): + x1, y1, x2, y2 = detections.xyxy[i].astype(int) + if self.color_by_track: + tracker_id = ( + detections.tracker_id[i] + if detections.tracker_id is not None + else None + ) + idx = tracker_id if tracker_id is not None else i + else: + class_id = ( + detections.class_id[i] if detections.class_id is not None else None + ) + idx = class_id if class_id is not None else i + color = ( + self.color.by_idx(idx) + if isinstance(self.color, ColorPalette) + else self.color + ) + + center = (int((x1 + x2) / 2), y2) + width = x2 - x1 + + cv2.ellipse( + scene, + center=center, + axes=(int(width), int(0.35 * width)), + angle=0.0, + startAngle=self.start_angle, + endAngle=self.end_angle, + color=color.as_bgr(), + thickness=self.thickness, + lineType=cv2.LINE_4, + ) + + return scene From 4b423fd4021c362e44b1058f0dc144809c7e982b Mon Sep 17 00:00:00 2001 From: hd Date: Fri, 7 Jul 2023 22:40:30 +0200 Subject: [PATCH 09/17] Fixed typo in annotator --- supervision/__init__.py | 4 ++-- supervision/annotators/core.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/supervision/__init__.py b/supervision/__init__.py index 21f3c81a2..28d29813d 100644 --- a/supervision/__init__.py +++ b/supervision/__init__.py @@ -9,8 +9,8 @@ from supervision.annotators.core import ( BoxAnnotator, BoxMaskAnnotator, - CorneredBoxAnotator, - EllipseAnotator, + CorneredBoxAnnotator, + EllipseAnnotator, LabelAnnotator, MaskAnnotator, PillowLabelAnnotator, diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index 94be97e61..7897fc465 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -574,7 +574,7 @@ def annotate( return scene -class CorneredBoxAnotator(BaseAnnotator): +class CorneredBoxAnnotator(BaseAnnotator): def __init__( self, color: Union[Color, ColorPalette] = ColorPalette.default(), @@ -699,7 +699,7 @@ def annotate( return scene -class EllipseAnotator(BaseAnnotator): +class EllipseAnnotator(BaseAnnotator): def __init__( self, color: Union[Color, ColorPalette] = ColorPalette.default(), From 9c4555a800a23b0bd3d2bada3583f32fed755ca6 Mon Sep 17 00:00:00 2001 From: hd Date: Wed, 19 Jul 2023 09:47:49 +0200 Subject: [PATCH 10/17] Added extra anchor points calculation --- supervision/annotators/core.py | 10 ++-------- supervision/detection/core.py | 30 ++++++++++++++++++++++++++++++ supervision/detection/track.py | 15 ++++----------- supervision/geometry/core.py | 7 +++++++ 4 files changed, 43 insertions(+), 19 deletions(-) diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index 7897fc465..2dd90590e 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -33,10 +33,7 @@ def __init__( self.color_by_track = color_by_track def annotate( - self, - scene: np.ndarray, - detections: Detections, - **kwargs + self, scene: np.ndarray, detections: Detections, **kwargs ) -> np.ndarray: """ Draws bounding boxes on the frame using the detections provided. @@ -113,10 +110,7 @@ def __init__( self.color_by_track = color_by_track def annotate( - self, - scene: np.ndarray, - detections: Detections, - **kwargs + self, scene: np.ndarray, detections: Detections, **kwargs ) -> np.ndarray: """ Overlays the masks on the given image based on the provided detections, with a specified opacity. diff --git a/supervision/detection/core.py b/supervision/detection/core.py index b423aaad7..5057374b1 100644 --- a/supervision/detection/core.py +++ b/supervision/detection/core.py @@ -441,11 +441,41 @@ def get_anchor_coordinates(self, anchor: Position) -> np.ndarray: (self.xyxy[:, 1] + self.xyxy[:, 3]) / 2, ] ).transpose() + elif anchor == Position.CENTER_LEFT: + return np.array( + [ + self.xyxy[:, 0], + (self.xyxy[:, 1] + self.xyxy[:, 3]) / 2, + ] + ).transpose() + + elif anchor == Position.CENTER_RIGHT: + return np.array( + [ + self.xyxy[:, 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() + elif anchor == Position.BOTTOM_LEFT: + return np.array([self.xyxy[:, 0], self.xyxy[:, 3]]).transpose() + + elif anchor == Position.BOTTOM_RIGHT: + return np.array([self.xyxy[:, 2], self.xyxy[:, 3]]).transpose() + elif anchor == Position.TOP_CENTER: + return np.array( + [(self.xyxy[:, 0] + self.xyxy[:, 2]) / 2, self.xyxy[:, 1]] + ).transpose() + elif anchor == Position.TOP_LEFT: + return np.array([self.xyxy[:, 0], self.xyxy[:, 1]]).transpose() + elif anchor == Position.TOP_RIGHT: + return np.array([self.xyxy[:, 2], self.xyxy[:, 1]]).transpose() + raise ValueError(f"{anchor} is not supported.") def __getitem__( diff --git a/supervision/detection/track.py b/supervision/detection/track.py index c0d5a6872..046ec12a3 100644 --- a/supervision/detection/track.py +++ b/supervision/detection/track.py @@ -38,19 +38,12 @@ def __init__( def update(self, frame_counter: int, detections: Detections) -> None: if detections.xyxy.shape[0] > 0: - xyxy = detections.xyxy + xy = detections.get_anchor_coordinates(anchor=self.position) - if self.position == Position.CENTER: - x = (xyxy[:, 0] + xyxy[:, 2]) / 2 - y = (xyxy[:, 1] + xyxy[:, 3]) / 2 - elif self.position == Position.BOTTOM_CENTER: - x = (xyxy[:, 0] + xyxy[:, 2]) / 2 - y = xyxy[:, 3] - - new_detections = np.zeros(shape=(xyxy.shape[0], 5)) + new_detections = np.zeros(shape=(xy.shape[0], 5)) new_detections[:, 0] = frame_counter - new_detections[:, 1] = x - new_detections[:, 2] = y + new_detections[:, 1] = xy[:, 0] + new_detections[:, 2] = xy[:, 1] new_detections[:, 3] = detections.class_id new_detections[:, 4] = detections.tracker_id self.storage = np.append(self.storage, new_detections, axis=0) diff --git a/supervision/geometry/core.py b/supervision/geometry/core.py index fc352f558..75f668674 100644 --- a/supervision/geometry/core.py +++ b/supervision/geometry/core.py @@ -7,7 +7,14 @@ class Position(Enum): CENTER = "CENTER" + CENTER_LEFT = "CENTER_LEFT" + CENTER_RIGHT = "CENTER_RIGHT" + TOP_CENTER = "TOP_CENTER" + TOP_LEFT = "TOP_LEFT" + TOP_RIGHT = "TOP_RIGHT" + BOTTOM_LEFT = "BOTTOM_LEFT" BOTTOM_CENTER = "BOTTOM_CENTER" + BOTTOM_RIGHT = "BOTTOM_RIGHT" @classmethod def list(cls): From c4d296fbc6d924876b9f255e3541e84dde7fbe1c Mon Sep 17 00:00:00 2001 From: hd Date: Wed, 19 Jul 2023 10:58:39 +0200 Subject: [PATCH 11/17] Ready for review --- supervision/annotators/core.py | 19 +++++++------ supervision/detection/track.py | 51 +++++++++++++++++++++++++--------- 2 files changed, 49 insertions(+), 21 deletions(-) diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index 2dd90590e..0305c67ad 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -416,13 +416,14 @@ def annotate( >>> track_annotator.annotate(scene) """ img_h, img_w, _ = scene.shape - unique_ids = np.unique(self.track_storage.storage[:, -1]) + unique_ids = np.unique(self.track_storage.tracker_id) for unique_id in unique_ids: - valid = np.where(self.track_storage.storage[:, -1] == unique_id)[0] - - frames = self.track_storage.storage[valid, 0] + valid = np.where(self.track_storage.tracker_id == unique_id)[0] + if valid.shape[0] == 0: + continue + frames = self.track_storage.frame_id[valid] latest_frame = np.argmax(frames) - points_to_draw = self.track_storage.storage[valid, 1:3] + points_to_draw = self.track_storage.xy[valid] n_pts = points_to_draw.shape[0] headx, heady = int(points_to_draw[latest_frame][0]), int( @@ -430,10 +431,12 @@ def annotate( ) if headx > self.boundry_tolerance and heady > self.boundry_tolerance: - if self.color_by_track: + idx = None + if not self.color_by_track and self.track_storage.class_id.shape[0] > 0: + class_id = self.track_storage.class_id[valid][0] + idx = int(class_id) + if self.color_by_track or idx is None: idx = int(unique_id) - else: - idx = int(self.track_storage.storage[0, -2]) color = ( self.color.by_idx(idx) if isinstance(self.color, ColorPalette) diff --git a/supervision/detection/track.py b/supervision/detection/track.py index 046ec12a3..8027c34b4 100644 --- a/supervision/detection/track.py +++ b/supervision/detection/track.py @@ -1,3 +1,5 @@ +from typing import Optional, Tuple + import numpy as np from supervision.detection.core import Detections @@ -34,28 +36,51 @@ def __init__( self.position = position self.max_length = max_length self.frame_counter = 0 - self.storage = np.zeros((0, 5)) + + self.xy = np.zeros((0, 2)) + self.confidence = np.zeros(0) + self.class_id = np.zeros(0) + self.tracker_id = np.zeros(0) + self.frame_id = np.zeros(0) def update(self, frame_counter: int, detections: Detections) -> None: - if detections.xyxy.shape[0] > 0: + n_dets = detections.xyxy.shape[0] + if n_dets > 0 and detections.tracker_id.shape[0] > 0: xy = detections.get_anchor_coordinates(anchor=self.position) - new_detections = np.zeros(shape=(xy.shape[0], 5)) - new_detections[:, 0] = frame_counter - new_detections[:, 1] = xy[:, 0] - new_detections[:, 2] = xy[:, 1] - new_detections[:, 3] = detections.class_id - new_detections[:, 4] = detections.tracker_id - self.storage = np.append(self.storage, new_detections, axis=0) + self.tracker_id = np.append(self.tracker_id, detections.tracker_id, axis=0) + + self.xy = np.append(self.xy, xy, axis=0) + if detections.class_id is not None: + self.class_id = np.append(self.class_id, detections.class_id, axis=0) + if detections.confidence is not None: + self.confidence = np.append( + self.confidence, detections.confidence, axis=0 + ) + self.frame_id = np.append( + self.frame_id, np.full((n_dets), fill_value=frame_counter), axis=0 + ) self.frame_counter = frame_counter self._remove_lost(tracker_ids=detections.tracker_id) self._remove_previous() def _remove_previous(self) -> None: to_remove_frames = self.frame_counter - self.max_length - if self.storage.shape[0] > 0: - self.storage = self.storage[self.storage[:, 0] > to_remove_frames] + if self.frame_id.shape[0] > 0: + valid = self.frame_id > to_remove_frames + self._remove(valid) def _remove_lost(self, tracker_ids) -> None: - if self.storage.shape[0] > 0: - self.storage = self.storage[np.isin(self.storage[:, -1], tracker_ids)] + if self.tracker_id.shape[0] > 0: + valid = np.isin(self.tracker_id, tracker_ids) + self._remove(valid) + + def _remove(self, valid: Optional[Tuple]) -> None: + if valid.size > 0: + self.xy = self.xy[valid] + self.frame_id = self.frame_id[valid] + self.tracker_id = self.tracker_id[valid] + if self.confidence.shape[0] > 0: + self.confidence = self.confidence[valid] + if self.class_id.shape[0] > 0: + self.class_id = self.class_id[valid] From f0cef99de73660f961bb2b7938befd4ca436584c Mon Sep 17 00:00:00 2001 From: hd Date: Fri, 21 Jul 2023 11:53:44 +0200 Subject: [PATCH 12/17] ready for review --- supervision/__init__.py | 10 +-- supervision/annotators/composable.py | 45 +--------- supervision/annotators/core.py | 120 ++++++++------------------- supervision/detection/track.py | 86 ------------------- 4 files changed, 38 insertions(+), 223 deletions(-) delete mode 100644 supervision/detection/track.py diff --git a/supervision/__init__.py b/supervision/__init__.py index 28d29813d..eb04df369 100644 --- a/supervision/__init__.py +++ b/supervision/__init__.py @@ -1,11 +1,6 @@ __version__ = "0.11.1" -from supervision.annotators.composable import ( - DetectionAnnotator, - SegmentationAnnotator, - TrackAnnotator, - TrackedDetectionAnnotator, -) +from supervision.annotators.composable import DetectionAnnotator, SegmentationAnnotator from supervision.annotators.core import ( BoxAnnotator, BoxMaskAnnotator, @@ -14,7 +9,7 @@ LabelAnnotator, MaskAnnotator, PillowLabelAnnotator, - TrackAnnotator, + build_label_formatter, ) from supervision.classification.core import Classifications from supervision.dataset.core import ( @@ -25,7 +20,6 @@ from supervision.detection.core import Detections from supervision.detection.line_counter import LineZone, LineZoneAnnotator from supervision.detection.tools.polygon_zone import PolygonZone, PolygonZoneAnnotator -from supervision.detection.track import TrackStorage from supervision.detection.utils import ( box_iou_batch, filter_polygons_by_area, diff --git a/supervision/annotators/composable.py b/supervision/annotators/composable.py index a3db9e7c1..5ee3cd31f 100644 --- a/supervision/annotators/composable.py +++ b/supervision/annotators/composable.py @@ -3,14 +3,8 @@ import numpy as np -from supervision.annotators.core import ( - BoxAnnotator, - LabelAnnotator, - MaskAnnotator, - TrackAnnotator, -) +from supervision.annotators.core import BoxAnnotator, LabelAnnotator, MaskAnnotator from supervision.detection.core import Detections -from supervision.detection.track import TrackStorage class ComposableAnnotator(ABC): @@ -32,8 +26,6 @@ def annotate( detections=detections, labels=labels, ) - elif isinstance(annotator, TrackAnnotator): - scene = annotator.annotate(scene=scene) else: scene = annotator.annotate(scene=scene, detections=detections) return scene @@ -100,38 +92,3 @@ def __init__( ] if not skip_label: self.annotators.append(LabelAnnotator(color_by_track=color_by_track)) - - -class TrackedDetectionAnnotator(ComposableAnnotator): - """ - A class for drawing object trajectories, bounding box and tracker ids on an image using provided detections. - User have freedom to choose color based on class_ids or tracker_ids - Example: - ```python - >>> import supervision as sv - - >>> classes = ['person', ...] - >>> image = ... - >>> detections = sv.Detections(...) - - >>> tracked_detection_annotator = sv.TrackedDetectionAnnotator() - >>> annotated_frame = tracked_detection_annotator.annotate( - ... scene=image.copy(), - ... detections=detections - ... ) - ``` - """ - - def __init__( - self, - tracks: TrackStorage, - color_by_track: bool = False, - skip_label: bool = False, - ): - super().__init__() - self.annotators = [ - BoxAnnotator(color_by_track=color_by_track), - TrackAnnotator(tracks=tracks, color_by_track=color_by_track), - ] - if not skip_label: - self.annotators.append(LabelAnnotator(color_by_track=color_by_track)) diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index 0305c67ad..1231bf07c 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -1,13 +1,12 @@ import os.path from abc import ABC, abstractmethod -from typing import List, Optional, Union +from typing import Callable, List, Optional, Union import cv2 import numpy as np from PIL import Image, ImageDraw, ImageFont from supervision.detection.core import Detections -from supervision.detection.track import TrackStorage from supervision.draw.color import Color, ColorPalette @@ -172,6 +171,19 @@ def annotate( return scene +def default_label_formatter( + detections: Detections, +) -> List[str]: + return [str(class_id) for class_id in detections.class_id] + + +def build_label_formatter(classes: List[str]) -> Callable[[Detections], List[str]]: + def default_label_formatter(detections: Detections) -> List[str]: + return [classes[class_id] for class_id in detections.class_id] + + return default_label_formatter + + class LabelAnnotator(BaseAnnotator): """ A class for putting text on an image using provided detections. @@ -189,20 +201,30 @@ def __init__( text_thickness: int = 1, text_padding: int = 10, color_by_track: bool = False, + classes: Optional[List[str]] = None, + label_formatter: Optional[ + Callable[[Detections], List[str]] + ] = default_label_formatter, ): + """ + Args: + color_by_track: pick color by tracker id + classes: Optional list of class name + label_formatter: Optional callback function for label generator, avoided if classes is provided + """ 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 self.color_by_track = color_by_track + self.classes = classes + self.label_formatter = label_formatter def annotate( self, scene: np.ndarray, detections: Detections, - labels: Optional[List[str]] = None, ) -> np.ndarray: """ Draws text on the frame using the detections provided and label. @@ -210,7 +232,6 @@ def annotate( Args: 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` are not provided, corresponding `class_id` will be used as label. Returns: np.ndarray: The image with the bounding boxes drawn on it @@ -231,12 +252,18 @@ def annotate( >>> annotated_frame = label_annotator.annotate( ... scene=image.copy(), ... detections=detections, - ... labels=labels ... ) ``` """ font = cv2.FONT_HERSHEY_SIMPLEX + labels = None + if self.classes: + label_formatter = build_label_formatter(self.classes) + labels = label_formatter(detections) + else: + labels = self.label_formatter(detections=detections) + for i in range(len(detections)): x1, y1, x2, y2 = detections.xyxy[i].astype(int) if self.color_by_track: @@ -374,85 +401,8 @@ def annotate( return scene -class TrackAnnotator(BaseAnnotator): - """ - Initialize TrackAnnotator - """ - - def __init__( - self, - tracks: TrackStorage, - color: Union[Color, ColorPalette] = ColorPalette.default(), - thickness: int = 2, - color_by_track: bool = False, - ): - self.track_storage = tracks - self.color: Union[Color, ColorPalette] = color - self.thickness: int = thickness - self.boundry_tolerance = 20 - self.color_by_track = color_by_track - - def annotate( - self, - scene: np.ndarray, - ) -> np.ndarray: - """ - Draws the object trajectory on the frame using the trace provided. - Attributes: - scene (np.ndarray): The image on which the object trajectories will be drawn. - tracks (TrackStorage): The track storage that will be used to draw the previous and current position. - - Returns: - np.ndarray: The image with the object trajectories on it. - ```python - >>> import supervision as sv - >>> track_storage = sv.TrackStorage() - >>> track_annotator = sv.TrackAnnotator(track_storage) - >>> for frame in sv.get_video_frames_generator(source_path='source_video.mp4'): - >>> detections = sv.Detections(...) - >>> tracked_objects = tracker(...) - >>> tracked_detections = sv.Detections(tracked_objects) - >>> track_storage.update(tracked_detections) - >>> track_annotator.annotate(scene) - """ - img_h, img_w, _ = scene.shape - unique_ids = np.unique(self.track_storage.tracker_id) - for unique_id in unique_ids: - valid = np.where(self.track_storage.tracker_id == unique_id)[0] - if valid.shape[0] == 0: - continue - frames = self.track_storage.frame_id[valid] - latest_frame = np.argmax(frames) - points_to_draw = self.track_storage.xy[valid] - - n_pts = points_to_draw.shape[0] - headx, heady = int(points_to_draw[latest_frame][0]), int( - points_to_draw[latest_frame][1] - ) - - if headx > self.boundry_tolerance and heady > self.boundry_tolerance: - idx = None - if not self.color_by_track and self.track_storage.class_id.shape[0] > 0: - class_id = self.track_storage.class_id[valid][0] - idx = int(class_id) - if self.color_by_track or idx is None: - idx = int(unique_id) - color = ( - self.color.by_idx(idx) - if isinstance(self.color, ColorPalette) - else self.color - ) - - for i in range(n_pts - 1): - px, py = int(points_to_draw[i][0]), int(points_to_draw[i][1]) - qx, qy = int(points_to_draw[i + 1][0]), int( - points_to_draw[i + 1][1] - ) - cv2.line(scene, (px, py), (qx, qy), color.as_bgr(), self.thickness) - scene = cv2.circle( - scene, (headx, heady), int(10), color.as_bgr(), thickness=-1 - ) - return scene +def default_label_formatter(detections: Detections) -> List[str]: + return [str(class_id) for class_id in detections.class_id] class PillowLabelAnnotator(BaseAnnotator): diff --git a/supervision/detection/track.py b/supervision/detection/track.py deleted file mode 100644 index 8027c34b4..000000000 --- a/supervision/detection/track.py +++ /dev/null @@ -1,86 +0,0 @@ -from typing import Optional, Tuple - -import numpy as np - -from supervision.detection.core import Detections -from supervision.geometry.core import Position - - -class TrackStorage: - """ - Trace the object trajectory with Detections with tracker id - """ - - def __init__( - self, - position: Position = Position.CENTER, - max_length: int = 10, - ): - """ - Initialize the track storage to store tracked detections from tracker - Args: - position [sv.Position]: Trace position whethere mid-point or bottom center point - max_length [int]: Length of previous detections to annotate - Detections objects are stored as numpy ndarray with #counter, x, y, class, track_id - - Example: - ```python - >>> import supervision as sv - >>> track_storage = sv.TrackStorage() - >>> for frame in sv.get_video_frames_generator(source_path='source_video.mp4'): - >>> detections = sv.Detections(...) - >>> tracked_objects = tracker(...) - >>> tracked_detections = sv.Detections(tracked_objects) - >>> track_storage.update(tracked_detections) - """ - self.position = position - self.max_length = max_length - self.frame_counter = 0 - - self.xy = np.zeros((0, 2)) - self.confidence = np.zeros(0) - self.class_id = np.zeros(0) - self.tracker_id = np.zeros(0) - self.frame_id = np.zeros(0) - - def update(self, frame_counter: int, detections: Detections) -> None: - n_dets = detections.xyxy.shape[0] - if n_dets > 0 and detections.tracker_id.shape[0] > 0: - xy = detections.get_anchor_coordinates(anchor=self.position) - - self.tracker_id = np.append(self.tracker_id, detections.tracker_id, axis=0) - - self.xy = np.append(self.xy, xy, axis=0) - if detections.class_id is not None: - self.class_id = np.append(self.class_id, detections.class_id, axis=0) - if detections.confidence is not None: - self.confidence = np.append( - self.confidence, detections.confidence, axis=0 - ) - self.frame_id = np.append( - self.frame_id, np.full((n_dets), fill_value=frame_counter), axis=0 - ) - self.frame_counter = frame_counter - self._remove_lost(tracker_ids=detections.tracker_id) - self._remove_previous() - - def _remove_previous(self) -> None: - to_remove_frames = self.frame_counter - self.max_length - if self.frame_id.shape[0] > 0: - valid = self.frame_id > to_remove_frames - self._remove(valid) - - def _remove_lost(self, tracker_ids) -> None: - if self.tracker_id.shape[0] > 0: - valid = np.isin(self.tracker_id, tracker_ids) - self._remove(valid) - - def _remove(self, valid: Optional[Tuple]) -> None: - if valid.size > 0: - self.xy = self.xy[valid] - self.frame_id = self.frame_id[valid] - self.tracker_id = self.tracker_id[valid] - if self.confidence.shape[0] > 0: - self.confidence = self.confidence[valid] - if self.class_id.shape[0] > 0: - self.class_id = self.class_id[valid] From ace6ce363f52af40f51d0b3c937f64212d058272 Mon Sep 17 00:00:00 2001 From: hd Date: Fri, 21 Jul 2023 12:54:19 +0200 Subject: [PATCH 13/17] Tests added for anchor fn --- test/detection/test_core.py | 48 ++++++++++++++++++++++++++++++++++++- 1 file changed, 47 insertions(+), 1 deletion(-) diff --git a/test/detection/test_core.py b/test/detection/test_core.py index 08b6a9d4b..afb315968 100644 --- a/test/detection/test_core.py +++ b/test/detection/test_core.py @@ -2,7 +2,7 @@ import pytest -from supervision import Detections +from supervision import Detections, Position from typing import Optional, Union, List @@ -243,3 +243,49 @@ def test_merge( with exception: result = Detections.merge(detections_list=detections_list) assert result == expected_result + + +def test_get_anchor_coordinates() -> None: + detections = mock_detections( + xyxy=[ + [10, 10, 20, 20], + [20, 20, 30, 30] + ] + ) + result = detections.get_anchor_coordinates(Position.CENTER).tolist() + expected_result = [[15, 15], [25, 25]] + assert result == expected_result + + result = detections.get_anchor_coordinates(Position.CENTER_LEFT).tolist() + expected_result = [[10, 15], [20, 25]] + assert result == expected_result + + result = detections.get_anchor_coordinates(Position.CENTER_RIGHT).tolist() + expected_result = [[20, 15], [30, 25]] + assert result == expected_result + + result = detections.get_anchor_coordinates(Position.TOP_CENTER).tolist() + expected_result = [[15, 10], [25, 20]] + assert result == expected_result + + result = detections.get_anchor_coordinates(Position.TOP_LEFT).tolist() + expected_result = [[10, 10], [20, 20]] + assert result == expected_result + + result = detections.get_anchor_coordinates(Position.TOP_RIGHT).tolist() + expected_result = [[20, 10], [30, 20]] + assert result == expected_result + + result = detections.get_anchor_coordinates(Position.BOTTOM_CENTER).tolist() + expected_result = [[15, 20], [25, 30]] + assert result == expected_result + + result = detections.get_anchor_coordinates(Position.BOTTOM_LEFT).tolist() + expected_result = [[10, 20], [20, 30]] + assert result == expected_result + + result = detections.get_anchor_coordinates(Position.BOTTOM_RIGHT).tolist() + expected_result = [[20, 20], [30, 30]] + assert result == expected_result + + From fc4bb27aa8a70df923a1fb0de5c26a45129c015e Mon Sep 17 00:00:00 2001 From: hd Date: Fri, 21 Jul 2023 12:54:44 +0200 Subject: [PATCH 14/17] removing extra spacing --- test/detection/test_core.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/test/detection/test_core.py b/test/detection/test_core.py index afb315968..b750b710a 100644 --- a/test/detection/test_core.py +++ b/test/detection/test_core.py @@ -287,5 +287,3 @@ def test_get_anchor_coordinates() -> None: result = detections.get_anchor_coordinates(Position.BOTTOM_RIGHT).tolist() expected_result = [[20, 20], [30, 30]] assert result == expected_result - - From 4cd9f05fab70e0b3f74c3a5e73d0128271fea1c1 Mon Sep 17 00:00:00 2001 From: SkalskiP Date: Sun, 23 Jul 2023 20:08:34 +0200 Subject: [PATCH 15/17] =?UTF-8?q?=F0=9F=A7=AA=20reformatted=20tests?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- supervision/detection/core.py | 4 -- test/dataset/test_core.py | 1 - test/detection/test_core.py | 122 ++++++++++++++++++++++------------ 3 files changed, 80 insertions(+), 47 deletions(-) diff --git a/supervision/detection/core.py b/supervision/detection/core.py index 5057374b1..ac007633b 100644 --- a/supervision/detection/core.py +++ b/supervision/detection/core.py @@ -448,7 +448,6 @@ def get_anchor_coordinates(self, anchor: Position) -> np.ndarray: (self.xyxy[:, 1] + self.xyxy[:, 3]) / 2, ] ).transpose() - elif anchor == Position.CENTER_RIGHT: return np.array( [ @@ -456,15 +455,12 @@ def get_anchor_coordinates(self, anchor: Position) -> np.ndarray: (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() - elif anchor == Position.BOTTOM_LEFT: return np.array([self.xyxy[:, 0], self.xyxy[:, 3]]).transpose() - elif anchor == Position.BOTTOM_RIGHT: return np.array([self.xyxy[:, 2], self.xyxy[:, 3]]).transpose() elif anchor == Position.TOP_CENTER: diff --git a/test/dataset/test_core.py b/test/dataset/test_core.py index a6252afce..356b6270f 100644 --- a/test/dataset/test_core.py +++ b/test/dataset/test_core.py @@ -172,5 +172,4 @@ def test_dataset_merge( ) -> None: with exception: result = DetectionDataset.merge(dataset_list=dataset_list) - print(result.images.keys()) assert result == expected_result diff --git a/test/detection/test_core.py b/test/detection/test_core.py index b750b710a..e03dd1b7c 100644 --- a/test/detection/test_core.py +++ b/test/detection/test_core.py @@ -245,45 +245,83 @@ def test_merge( assert result == expected_result -def test_get_anchor_coordinates() -> None: - detections = mock_detections( - xyxy=[ - [10, 10, 20, 20], - [20, 20, 30, 30] - ] - ) - result = detections.get_anchor_coordinates(Position.CENTER).tolist() - expected_result = [[15, 15], [25, 25]] - assert result == expected_result - - result = detections.get_anchor_coordinates(Position.CENTER_LEFT).tolist() - expected_result = [[10, 15], [20, 25]] - assert result == expected_result - - result = detections.get_anchor_coordinates(Position.CENTER_RIGHT).tolist() - expected_result = [[20, 15], [30, 25]] - assert result == expected_result - - result = detections.get_anchor_coordinates(Position.TOP_CENTER).tolist() - expected_result = [[15, 10], [25, 20]] - assert result == expected_result - - result = detections.get_anchor_coordinates(Position.TOP_LEFT).tolist() - expected_result = [[10, 10], [20, 20]] - assert result == expected_result - - result = detections.get_anchor_coordinates(Position.TOP_RIGHT).tolist() - expected_result = [[20, 10], [30, 20]] - assert result == expected_result - - result = detections.get_anchor_coordinates(Position.BOTTOM_CENTER).tolist() - expected_result = [[15, 20], [25, 30]] - assert result == expected_result - - result = detections.get_anchor_coordinates(Position.BOTTOM_LEFT).tolist() - expected_result = [[10, 20], [20, 30]] - assert result == expected_result - - result = detections.get_anchor_coordinates(Position.BOTTOM_RIGHT).tolist() - expected_result = [[20, 20], [30, 30]] - assert result == expected_result +@pytest.mark.parametrize( + 'detections, anchor, expected_result, exception', + [ + ( + Detections.empty(), + Position.CENTER, + np.empty((0, 2), dtype=np.float32), + DoesNotRaise() + ), # empty detections + ( + mock_detections(xyxy=[[10, 10, 20, 20]]), + Position.CENTER, + np.array([[15, 15]], dtype=np.float32), + DoesNotRaise() + ), # single detection; center anchor + ( + mock_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]]), + Position.CENTER, + np.array([[15, 15], [25, 25]], dtype=np.float32), + DoesNotRaise() + ), # two detections; center anchor + ( + mock_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]]), + Position.CENTER_LEFT, + np.array([[10, 15], [20, 25]], dtype=np.float32), + DoesNotRaise() + ), # two detections; center left anchor + ( + mock_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]]), + Position.CENTER_RIGHT, + np.array([[20, 15], [30, 25]], dtype=np.float32), + DoesNotRaise() + ), # two detections; center right anchor + ( + mock_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]]), + Position.TOP_CENTER, + np.array([[15, 10], [25, 20]], dtype=np.float32), + DoesNotRaise() + ), # two detections; top center anchor + ( + mock_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]]), + Position.TOP_LEFT, + np.array([[10, 10], [20, 20]], dtype=np.float32), + DoesNotRaise() + ), # two detections; top left anchor + ( + mock_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]]), + Position.TOP_RIGHT, + np.array([[20, 10], [30, 20]], dtype=np.float32), + DoesNotRaise() + ), # two detections; top right anchor + ( + mock_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]]), + Position.BOTTOM_CENTER, + np.array([[15, 20], [25, 30]], dtype=np.float32), + DoesNotRaise() + ), # two detections; bottom center anchor + ( + mock_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]]), + Position.BOTTOM_LEFT, + np.array([[10, 20], [20, 30]], dtype=np.float32), + DoesNotRaise() + ), # two detections; bottom left anchor + ( + mock_detections(xyxy=[[10, 10, 20, 20], [20, 20, 30, 30]]), + Position.BOTTOM_RIGHT, + np.array([[20, 20], [30, 30]], dtype=np.float32), + DoesNotRaise() + ), # two detections; bottom right anchor + ] +) +def test_get_anchor_coordinates( + detections: Detections, + anchor: Position, + expected_result: np.ndarray, + exception: Exception +) -> None: + result = detections.get_anchor_coordinates(anchor) + with exception: + assert np.array_equal(result, expected_result) From ba03fefbd8a9eeafc37a99fc5721ce5866d9af65 Mon Sep 17 00:00:00 2001 From: SkalskiP Date: Sun, 23 Jul 2023 20:39:12 +0200 Subject: [PATCH 16/17] =?UTF-8?q?=F0=9F=A7=B9=E2=80=8D=F0=9F=92=A8=20clean?= =?UTF-8?q?up=20in=20progress?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- supervision/__init__.py | 7 ++++--- supervision/annotators/core.py | 19 +++++++++++-------- test/detection/test_core.py | 3 ++- 3 files changed, 17 insertions(+), 12 deletions(-) diff --git a/supervision/__init__.py b/supervision/__init__.py index eb04df369..6f7caf953 100644 --- a/supervision/__init__.py +++ b/supervision/__init__.py @@ -2,15 +2,16 @@ from supervision.annotators.composable import DetectionAnnotator, SegmentationAnnotator from supervision.annotators.core import ( - BoxAnnotator, + BoxLineAnnotator, BoxMaskAnnotator, - CorneredBoxAnnotator, + BoxCornerAnnotator, EllipseAnnotator, LabelAnnotator, MaskAnnotator, - PillowLabelAnnotator, + LabelAdvancedAnnotator, build_label_formatter, ) +from supervision.detection.annotate import BoxAnnotator, MaskAnnotator from supervision.classification.core import Classifications from supervision.dataset.core import ( BaseDataset, diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index 1231bf07c..78be2be5e 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -15,8 +15,12 @@ class BaseAnnotator(ABC): def annotate(self, scene: np.ndarray, detections: Detections) -> np.ndarray: pass + @staticmethod + def resolve_annotation_color(color: Union[Color, ColorPalette], by_track: bool, detections: Detections) -> Color: + pass + -class BoxAnnotator(BaseAnnotator): +class BoxLineAnnotator(BaseAnnotator): """ Basic bounding box annotation class """ @@ -41,18 +45,17 @@ def annotate( 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 Returns: - np.ndarray: The image with the bounding boxes drawn on it + np.ndarray: The image with the bounding boxes drawn on it. Example: ```python >>> import supervision as sv - >>> classes = ['person', ...] >>> image = ... >>> detections = sv.Detections(...) - >>> box_annotator = sv.BoxAnnotator() - >>> annotated_frame = box_annotator.annotate( + >>> box_line_annotator = sv.BoxLineAnnotator() + >>> annotated_frame = box_line_annotator.annotate( ... scene=image.copy(), ... detections=detections ... ) @@ -405,7 +408,7 @@ def default_label_formatter(detections: Detections) -> List[str]: return [str(class_id) for class_id in detections.class_id] -class PillowLabelAnnotator(BaseAnnotator): +class LabelAdvancedAnnotator(BaseAnnotator): def __init__( self, color: Union[Color, ColorPalette] = ColorPalette.default(), @@ -448,7 +451,7 @@ def annotate( >>> image = ... >>> detections = sv.Detections(...) - >>> pil_label_annotator = sv.PillowLabelAnnotator() + >>> pil_label_annotator = sv.LabelAdvancedAnnotator() >>> labels = [ ... f"{classes[class_id]} {confidence:0.2f}" ... for _, _, confidence, class_id, _ @@ -521,7 +524,7 @@ def annotate( return scene -class CorneredBoxAnnotator(BaseAnnotator): +class BoxCornerAnnotator(BaseAnnotator): def __init__( self, color: Union[Color, ColorPalette] = ColorPalette.default(), diff --git a/test/detection/test_core.py b/test/detection/test_core.py index e03dd1b7c..01e6b9f87 100644 --- a/test/detection/test_core.py +++ b/test/detection/test_core.py @@ -2,7 +2,8 @@ import pytest -from supervision import Detections, Position +from supervision.detection.core import Detections +from supervision.geometry.core import Position from typing import Optional, Union, List From cc2e08ac7229bc6aac329c32c977ba3f93452e90 Mon Sep 17 00:00:00 2001 From: SkalskiP Date: Sun, 23 Jul 2023 20:40:46 +0200 Subject: [PATCH 17/17] =?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 --- supervision/__init__.py | 6 +++--- supervision/annotators/core.py | 4 +++- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/supervision/__init__.py b/supervision/__init__.py index 6f7caf953..d446c1655 100644 --- a/supervision/__init__.py +++ b/supervision/__init__.py @@ -2,22 +2,22 @@ from supervision.annotators.composable import DetectionAnnotator, SegmentationAnnotator from supervision.annotators.core import ( + BoxCornerAnnotator, BoxLineAnnotator, BoxMaskAnnotator, - BoxCornerAnnotator, EllipseAnnotator, + LabelAdvancedAnnotator, LabelAnnotator, MaskAnnotator, - LabelAdvancedAnnotator, build_label_formatter, ) -from supervision.detection.annotate import BoxAnnotator, MaskAnnotator from supervision.classification.core import Classifications from supervision.dataset.core import ( BaseDataset, ClassificationDataset, DetectionDataset, ) +from supervision.detection.annotate import BoxAnnotator, MaskAnnotator from supervision.detection.core import Detections from supervision.detection.line_counter import LineZone, LineZoneAnnotator from supervision.detection.tools.polygon_zone import PolygonZone, PolygonZoneAnnotator diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index 78be2be5e..6b72a065c 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -16,7 +16,9 @@ def annotate(self, scene: np.ndarray, detections: Detections) -> np.ndarray: pass @staticmethod - def resolve_annotation_color(color: Union[Color, ColorPalette], by_track: bool, detections: Detections) -> Color: + def resolve_annotation_color( + color: Union[Color, ColorPalette], by_track: bool, detections: Detections + ) -> Color: pass