diff --git a/setup.py b/setup.py index ac3add4f..69ca93b6 100644 --- a/setup.py +++ b/setup.py @@ -27,7 +27,8 @@ setuptools.setup( '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 30694450..d446c165 100644 --- a/supervision/__init__.py +++ b/supervision/__init__.py @@ -1,5 +1,16 @@ __version__ = "0.11.1" +from supervision.annotators.composable import DetectionAnnotator, SegmentationAnnotator +from supervision.annotators.core import ( + BoxCornerAnnotator, + BoxLineAnnotator, + BoxMaskAnnotator, + EllipseAnnotator, + LabelAdvancedAnnotator, + LabelAnnotator, + MaskAnnotator, + build_label_formatter, +) from supervision.classification.core import Classifications from supervision.dataset.core import ( BaseDataset, diff --git a/supervision/annotators/composable.py b/supervision/annotators/composable.py index 4c0eaeb9..5ee3cd31 100644 --- a/supervision/annotators/composable.py +++ b/supervision/annotators/composable.py @@ -1,18 +1,94 @@ 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 BoxAnnotator, LabelAnnotator, MaskAnnotator 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, + ) -> 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, + ) + else: + scene = annotator.annotate(scene=scene, detections=detections) + 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, + color_by_track: bool = False, + skip_label: bool = False, + ): + super().__init__() + 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): + """ + 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, + color_by_track: bool = False, + skip_label: bool = False, + ): + super().__init__() + self.annotators = [ + 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)) diff --git a/supervision/annotators/core.py b/supervision/annotators/core.py index 9e477a1a..6b72a065 100644 --- a/supervision/annotators/core.py +++ b/supervision/annotators/core.py @@ -1,7 +1,10 @@ +import os.path from abc import ABC, abstractmethod -from typing import 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.draw.color import Color, ColorPalette @@ -12,12 +15,85 @@ class BaseAnnotator(ABC): def annotate(self, scene: np.ndarray, detections: Detections) -> np.ndarray: pass - -class BoxAnnotator(BaseAnnotator): - def annotate(self, scene: np.ndarray, detections: Detections) -> np.ndarray: + @staticmethod + def resolve_annotation_color( + color: Union[Color, ColorPalette], by_track: bool, detections: Detections + ) -> Color: pass +class BoxLineAnnotator(BaseAnnotator): + """ + Basic bounding box annotation class + """ + + 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, **kwargs + ) -> 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 + Returns: + np.ndarray: The image with the bounding boxes drawn on it. + + Example: + ```python + >>> import supervision as sv + + >>> image = ... + >>> detections = sv.Detections(...) + + >>> box_line_annotator = sv.BoxLineAnnotator() + >>> annotated_frame = box_line_annotator.annotate( + ... scene=image.copy(), + ... detections=detections + ... ) + ``` + """ + 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 + ) + cv2.rectangle( + img=scene, + pt1=(x1, y1), + pt2=(x2, y2), + color=color.as_bgr(), + thickness=self.thickness, + ) + + return scene + + class MaskAnnotator(BaseAnnotator): """ A class for overlaying masks on an image using detections provided. @@ -31,11 +107,15 @@ class MaskAnnotator(BaseAnnotator): 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) -> np.ndarray: + 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. @@ -45,15 +125,38 @@ class MaskAnnotator(BaseAnnotator): 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 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 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) @@ -73,21 +176,556 @@ class MaskAnnotator(BaseAnnotator): 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): - def annotate(self, scene: np.ndarray, detections: Detections) -> np.ndarray: - pass + """ + 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 + """ -class PillowLabelAnnotator(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, + 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.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, + ) -> np.ndarray: + """ + Draws text on the frame using the detections provided and label. -class TrackAnnotator(BaseAnnotator): - def annotate(self, scene: np.ndarray, detections: Detections) -> np.ndarray: - pass + 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(...) + + >>> 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, + ... ) + ``` + """ + 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: + 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_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 BoxMaskAnnotator(BaseAnnotator): - def annotate(self, scene: np.ndarray, detections: Detections) -> np.ndarray: - pass + 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, + ) -> 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 + 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 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 + ) + 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 + + +def default_label_formatter(detections: Detections) -> List[str]: + return [str(class_id) for class_id in detections.class_id] + + +class LabelAdvancedAnnotator(BaseAnnotator): + def __init__( + self, + 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, + ): + 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, + ) -> 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. + 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.LabelAdvancedAnnotator() + >>> 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, + ... ) + ``` + """ + 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 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 + ) + + 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 BoxCornerAnnotator(BaseAnnotator): + 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, + ): + """ + 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 + 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 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 + ) + + 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 + + +class EllipseAnnotator(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 diff --git a/supervision/detection/core.py b/supervision/detection/core.py index b423aaad..ac007633 100644 --- a/supervision/detection/core.py +++ b/supervision/detection/core.py @@ -441,10 +441,36 @@ class Detections: (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.") diff --git a/supervision/geometry/core.py b/supervision/geometry/core.py index fc352f55..75f66867 100644 --- a/supervision/geometry/core.py +++ b/supervision/geometry/core.py @@ -7,7 +7,14 @@ from typing import Tuple 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): diff --git a/test/dataset/test_core.py b/test/dataset/test_core.py index a6252afc..356b6270 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 08b6a9d4..01e6b9f8 100644 --- a/test/detection/test_core.py +++ b/test/detection/test_core.py @@ -2,7 +2,8 @@ from contextlib import ExitStack as DoesNotRaise import pytest -from supervision import Detections +from supervision.detection.core import Detections +from supervision.geometry.core import Position from typing import Optional, Union, List @@ -243,3 +244,85 @@ def test_merge( with exception: result = Detections.merge(detections_list=detections_list) 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)