Skip to content

Annotate

BoxAnnotator

A class for drawing bounding boxes on an image using detections provided.

Attributes:

Name Type Description
color Union[Color, ColorPalette]

The color to draw the bounding box, can be a single color or a color palette

thickness int

The thickness of the bounding box lines, default is 2

text_color Color

The color of the text on the bounding box, default is white

text_scale float

The scale of the text on the bounding box, default is 0.5

text_thickness int

The thickness of the text on the bounding box, default is 1

text_padding int

The padding around the text on the bounding box, default is 5

Source code in supervision/detection/annotate.py
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
class BoxAnnotator:
    """
    A class for drawing bounding boxes on an image using detections provided.

    Attributes:
        color (Union[Color, ColorPalette]): The color to draw the bounding box, can be a single color or a color palette
        thickness (int): The thickness of the bounding box lines, default is 2
        text_color (Color): The color of the text on the bounding box, default is white
        text_scale (float): The scale of the text on the bounding box, default is 0.5
        text_thickness (int): The thickness of the text on the bounding box, default is 1
        text_padding (int): The padding around the text on the bounding box, default is 5

    """

    def __init__(
        self,
        color: Union[Color, ColorPalette] = ColorPalette.default(),
        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,
    ) -> np.ndarray:
        """
        Draws bounding boxes on the frame using the detections provided.

        Parameters:
            scene (np.ndarray): The image on which the bounding boxes will be drawn
            detections (Detections): The detections for which the bounding boxes will be drawn
            labels (Optional[List[str]]): An optional list of labels corresponding to each detection. If labels is provided, the confidence score of the detection will be replaced with the label.
            skip_label (bool): Is set to True, skips bounding box label annotation.
        Returns:
            np.ndarray: The image with the bounding boxes drawn on it
        """
        font = cv2.FONT_HERSHEY_SIMPLEX
        for i, (xyxy, confidence, class_id, tracker_id) in enumerate(detections):
            x1, y1, x2, y2 = xyxy.astype(int)
            idx = class_id if 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,
            )
            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

annotate(scene, detections, labels=None, skip_label=False)

Draws bounding boxes on the frame using the detections provided.

Parameters:

Name Type Description Default
scene ndarray

The image on which the bounding boxes will be drawn

required
detections Detections

The detections for which the bounding boxes will be drawn

required
labels Optional[List[str]]

An optional list of labels corresponding to each detection. If labels is provided, the confidence score of the detection will be replaced with the label.

None
skip_label bool

Is set to True, skips bounding box label annotation.

False

Returns: np.ndarray: The image with the bounding boxes drawn on it

Source code in supervision/detection/annotate.py
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
def annotate(
    self,
    scene: np.ndarray,
    detections: Detections,
    labels: Optional[List[str]] = None,
    skip_label: bool = False,
) -> np.ndarray:
    """
    Draws bounding boxes on the frame using the detections provided.

    Parameters:
        scene (np.ndarray): The image on which the bounding boxes will be drawn
        detections (Detections): The detections for which the bounding boxes will be drawn
        labels (Optional[List[str]]): An optional list of labels corresponding to each detection. If labels is provided, the confidence score of the detection will be replaced with the label.
        skip_label (bool): Is set to True, skips bounding box label annotation.
    Returns:
        np.ndarray: The image with the bounding boxes drawn on it
    """
    font = cv2.FONT_HERSHEY_SIMPLEX
    for i, (xyxy, confidence, class_id, tracker_id) in enumerate(detections):
        x1, y1, x2, y2 = xyxy.astype(int)
        idx = class_id if 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,
        )
        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