diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index 4ffde6cd..39a22fb7 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -68,7 +68,7 @@ jobs: cache: pip - name: Install depthai - run: pip install --extra-index-url https://artifacts.luxonis.com/artifactory/luxonis-python-release-local/ depthai==3.0.0a11 + run: pip install --extra-index-url https://artifacts.luxonis.com/artifactory/luxonis-python-release-local/ depthai==3.0.0a13 - name: Install package run: pip install -e .[dev] diff --git a/depthai_nodes/utils/annotation_helper.py b/depthai_nodes/utils/annotation_helper.py new file mode 100644 index 00000000..d248ca18 --- /dev/null +++ b/depthai_nodes/utils/annotation_helper.py @@ -0,0 +1,301 @@ +from datetime import timedelta +from typing import List, Optional, Tuple + +import depthai as dai +import numpy as np + +Point = Tuple[float, float] +ColorRGBA = Tuple[float, float, float, float] + + +class AnnotationHelper: + """Simplifies `dai.ImgAnnotation` creation. + + After calling the desired drawing methods, call the `build` method to create the `ImgAnnotations` message. + """ + + def __init__(self): + self.annotation: dai.ImgAnnotation = dai.ImgAnnotation() + + def draw_line( + self, pt1: Point, pt2: Point, color: ColorRGBA, thickness: float + ) -> "AnnotationHelper": + """Draws a line between two points. + + @param pt1: Start of the line + @type pt1: Point + @param pt2: End of the line + @type pt2: Point + @param color: Line color + @type color: ColorRGBA + @param thickness: Line thickness + @type thickness: float + @return: self + @rtype: AnnotationHelper + """ + line = dai.PointsAnnotation() + c = self._create_color(color) + line.fillColor = c + line.outlineColor = c + line.thickness = thickness + line.type = dai.PointsAnnotationType.LINE_STRIP + line.points = self._create_points_vector([pt1, pt2]) + self.annotation.points.append(line) + return self + + def draw_polyline( + self, + points: List[Point], + outline_color: ColorRGBA, + fill_color: Optional[ColorRGBA] = None, + thickness: float = 1, + closed: bool = False, + ) -> "AnnotationHelper": + """Draws a polyline. + + @param points: List of points of the polyline + @type points: list[Point] + @param outline_color: Outline color + @type outline_color: ColorRGBA + @param fill_color: Fill color (None for no fill), defaults to None + @type fill_color: ColorRGBA | None, optional + @param thickness: Line thickness, defaults to 1 + @type thickness: float, optional + @param closed: Creates polygon, instead of polyline if True, defaults to False + @type closed: bool, optional + @return: self + @rtype: AnnotationHelper + """ + points_type = ( + dai.PointsAnnotationType.LINE_STRIP + if not closed + else dai.PointsAnnotationType.LINE_LOOP + ) + points_annot = self._create_points_annotation( + points, outline_color, fill_color, points_type + ) + points_annot.thickness = thickness + self.annotation.points.append(points_annot) + return self + + def draw_points( + self, points: List[Point], color: ColorRGBA, thickness: float = 2 + ) -> "AnnotationHelper": + """Draws points. + + @param points: List of points to draw + @type points: list[Point] + @param color: Color of the points + @type color: ColorRGBA + @param thickness: Size of the points, defaults to 2 + @type thickness: float, optional + @return: self + @rtype: AnnotationHelper + """ + # TODO: Visualizer currently does not show dai.PointsAnnotationType.POINTS + points_annot = self._create_points_annotation( + points, color, None, dai.PointsAnnotationType.POINTS + ) + points_annot.thickness = thickness + self.annotation.points.append(points_annot) + return self + + def draw_circle( + self, + center: Point, + radius: float, + outline_color: ColorRGBA, + fill_color: Optional[ColorRGBA] = None, + thickness: float = 1, + ) -> "AnnotationHelper": + """Draws a circle. + + @param center: Center of the circle + @type center: Point + @param radius: Radius of the circle + @type radius: float + @param outline_color: Outline color + @type outline_color: ColorRGBA + @param fill_color: Fill color (None for no fill), defaults to None + @type fill_color: ColorRGBA | None, optional + @param thickness: Outline thickness, defaults to 1 + @type thickness: float, optional + @return: self + @rtype: AnnotationHelper + """ + circle = dai.CircleAnnotation() + circle.outlineColor = self._create_color(outline_color) + if fill_color is not None: + circle.fillColor = self._create_color(fill_color) + circle.thickness = thickness + circle.diameter = radius * 2 + circle.position = dai.Point2f(center[0], center[1]) + self.annotation.circles.append(circle) + return self + + def draw_rectangle( + self, + top_left: Point, + bottom_right: Point, + outline_color: ColorRGBA, + fill_color: Optional[ColorRGBA] = None, + thickness: float = 1, + ) -> "AnnotationHelper": + """Draws a rectangle. + + @param top_left: Top left corner of the rectangle + @type top_left: Point + @param bottom_right: Bottom right corner of the rectangle + @type bottom_right: Point + @param outline_color: Outline color + @type outline_color: ColorRGBA + @param fill_color: Fill color (None for no fill), defaults to None + @type fill_color: ColorRGBA | None, optional + @param thickness: Outline thickness, defaults to 1 + @type thickness: float, optional + @return: self + @rtype: AnnotationHelper + """ + points = [ + top_left, + (bottom_right[0], top_left[1]), + bottom_right, + (top_left[0], bottom_right[1]), + ] + self.draw_polyline(points, outline_color, fill_color, thickness, closed=True) + return self + + def draw_text( + self, + text: str, + position: Point, + color: ColorRGBA, + background_color: Optional[ColorRGBA] = None, + size: float = 32, + ) -> "AnnotationHelper": + """Draws text. + + @param text: Text string + @type text: str + @param position: Text position + @type position: Point + @param color: Text color + @type color: ColorRGBA + @param background_color: Background color (None for no background), defaults to + None + @type background_color: ColorRGBA | None, optional + @param size: Text size, defaults to 32 + @type size: float, optional + @return: self + @rtype: AnnotationHelper + """ + text_annot = dai.TextAnnotation() + text_annot.position = dai.Point2f(position[0], position[1]) + text_annot.text = text + text_annot.textColor = self._create_color(color) + text_annot.fontSize = size + if background_color is not None: + text_annot.backgroundColor = self._create_color(background_color) + self.annotation.texts.append(text_annot) + return self + + def draw_rotated_rect( + self, + center: Point, + size: Tuple[float, float], + angle: float, + outline_color: ColorRGBA, + fill_color: Optional[ColorRGBA] = None, + thickness: float = 1, + ) -> "AnnotationHelper": + """Draws a rotated rectangle. + + @param center: Center of the rectangle + @type center: Point + @param size: Size of the rectangle (width, height) + @type size: tuple[float, float] + @param angle: Angle of rotation in degrees + @type angle: float + @param outline_color: Outline color + @type outline_color: ColorRGBA + @param fill_color: Fill color (None for no fill), defaults to None + @type fill_color: ColorRGBA | None, optional + @param thickness: Outline thickness, defaults to 1 + @type thickness: float, optional + @return: self + @rtype: AnnotationHelper + """ + points = self._get_rotated_rect_points(center, size, angle) + self.draw_polyline(points, outline_color, fill_color, thickness, True) + return self + + def build(self, timestamp: timedelta, sequence_num: int) -> dai.ImgAnnotations: + """Creates an ImgAnnotations message. + + @param timestamp: Message timestamp + @type timestamp: timedelta + @param sequence_num: Message sequence number + @type sequence_num: int + @return: Created ImgAnnotations message + @rtype: dai.ImgAnnotations + """ + annotations_msg = dai.ImgAnnotations() + annotations_msg.annotations = dai.VectorImgAnnotation([self.annotation]) + annotations_msg.setTimestamp(timestamp) + annotations_msg.setSequenceNum(sequence_num) + return annotations_msg + + def _create_points_annotation( + self, + points: List[Point], + color: ColorRGBA, + fill_color: Optional[ColorRGBA], + type: dai.PointsAnnotationType, + ) -> dai.PointsAnnotation: + points_annot = dai.PointsAnnotation() + points_annot.outlineColor = self._create_color(color) + if fill_color is not None: + points_annot.fillColor = self._create_color(fill_color) + points_annot.type = type + points_annot.points = self._create_points_vector(points) + return points_annot + + def _create_color(self, color: ColorRGBA) -> dai.Color: + c = dai.Color() + c.a = color[3] + c.r = color[0] + c.g = color[1] + c.b = color[2] + return c + + def _get_rotated_rect_points( + self, center: Point, size: Tuple[float, float], angle: float + ) -> List[Point]: + cx, cy = center + width, height = size + angle_rad = np.radians(angle) + + # Half-dimensions + dx = width / 2 + dy = height / 2 + + # Define the corners relative to the center + corners = np.array([[-dx, -dy], [dx, -dy], [dx, dy], [-dx, dy]]) + + # Rotation matrix + rotation_matrix = np.array( + [ + [np.cos(angle_rad), -np.sin(angle_rad)], + [np.sin(angle_rad), np.cos(angle_rad)], + ] + ) + + # Rotate and translate the corners + rotated_corners = corners @ rotation_matrix.T + translated_corners = rotated_corners + np.array([cx, cy]) + + # Convert to list of tuples + return translated_corners.tolist() + + def _create_points_vector(self, points: List[Point]) -> dai.VectorPoint2f: + return dai.VectorPoint2f([dai.Point2f(pt[0], pt[1]) for pt in points]) diff --git a/tests/unittests/test_helpers/test_annotation_helper.py b/tests/unittests/test_helpers/test_annotation_helper.py new file mode 100644 index 00000000..06098f63 --- /dev/null +++ b/tests/unittests/test_helpers/test_annotation_helper.py @@ -0,0 +1,170 @@ +from datetime import timedelta + +import depthai as dai + +from depthai_nodes.utils.annotation_helper import AnnotationHelper + + +def test_empty_build(): + annotation_helper = AnnotationHelper() + annots = annotation_helper.build(timedelta(), 0) + assert len(annots.annotations[0].points) == 0 + assert len(annots.annotations[0].circles) == 0 + assert len(annots.annotations[0].texts) == 0 + + +def test_draw_line(): + annotation_helper = AnnotationHelper() + annotation_helper.draw_line((0.0, 0.0), (1.0, 1.0), (1.0, 0.0, 0.0, 1.0), 2) + annots = annotation_helper.build(timedelta(), 0) + + assert [(i.x, i.y) for i in annots.annotations[0].points[0].points] == [ + (0.0, 0.0), + (1.0, 1.0), + ] + assert annots.annotations[0].points[0].thickness == 2 + assert annots.annotations[0].points[0].type == dai.PointsAnnotationType.LINE_STRIP + assert annots.annotations[0].points[0].outlineColor.r == 1.0 + assert annots.annotations[0].points[0].outlineColor.g == 0.0 + assert annots.annotations[0].points[0].outlineColor.b == 0.0 + assert annots.annotations[0].points[0].outlineColor.a == 1.0 + + +def test_draw_polyline(): + annotation_helper = AnnotationHelper() + annotation_helper.draw_polyline( + [(0.0, 0.0), (0.5, 0.5), (1.0, 1.0)], + (1.0, 0.0, 0.0, 1.0), + (0.0, 1.0, 0.0, 1.0), + 2, + True, + ) + annots = annotation_helper.build(timedelta(), 0) + + assert [(i.x, i.y) for i in annots.annotations[0].points[0].points] == [ + (0.0, 0.0), + (0.5, 0.5), + (1.0, 1.0), + ] + assert annots.annotations[0].points[0].thickness == 2 + assert annots.annotations[0].points[0].type == dai.PointsAnnotationType.LINE_LOOP + assert annots.annotations[0].points[0].outlineColor.r == 1.0 + assert annots.annotations[0].points[0].outlineColor.g == 0.0 + assert annots.annotations[0].points[0].outlineColor.b == 0.0 + assert annots.annotations[0].points[0].outlineColor.a == 1.0 + assert annots.annotations[0].points[0].fillColor.r == 0.0 + assert annots.annotations[0].points[0].fillColor.g == 1.0 + assert annots.annotations[0].points[0].fillColor.b == 0.0 + assert annots.annotations[0].points[0].fillColor.a == 1.0 + + +def test_draw_polyline_closed(): + annotation_helper = AnnotationHelper() + annotation_helper.draw_polyline( + [(0.0, 0.0), (0.5, 0.5), (1.0, 1.0)], + (1.0, 0.0, 0.0, 1.0), + None, + 2, + False, + ) + annots = annotation_helper.build(timedelta(), 0) + + assert [(i.x, i.y) for i in annots.annotations[0].points[0].points] == [ + (0.0, 0.0), + (0.5, 0.5), + (1.0, 1.0), + ] + assert annots.annotations[0].points[0].thickness == 2 + assert annots.annotations[0].points[0].type == dai.PointsAnnotationType.LINE_STRIP + assert annots.annotations[0].points[0].outlineColor.r == 1.0 + assert annots.annotations[0].points[0].outlineColor.g == 0.0 + assert annots.annotations[0].points[0].outlineColor.b == 0.0 + assert annots.annotations[0].points[0].outlineColor.a == 1.0 + assert annots.annotations[0].points[0].fillColor.r == 0.0 + assert annots.annotations[0].points[0].fillColor.g == 0.0 + assert annots.annotations[0].points[0].fillColor.b == 0.0 + assert annots.annotations[0].points[0].fillColor.a == 0.0 + + +def test_draw_circle(): + annotation_helper = AnnotationHelper() + annotation_helper.draw_circle( + (0.0, 0.0), 0.5, (1.0, 0.0, 0.0, 1.0), (0.0, 1.0, 0.0, 1.0), 2 + ) + annots = annotation_helper.build(timedelta(), 0) + + assert annots.annotations[0].circles[0].position.x == 0.0 + assert annots.annotations[0].circles[0].position.y == 0.0 + assert annots.annotations[0].circles[0].diameter == 1.0 + assert annots.annotations[0].circles[0].outlineColor.r == 1.0 + assert annots.annotations[0].circles[0].outlineColor.g == 0.0 + assert annots.annotations[0].circles[0].outlineColor.b == 0.0 + assert annots.annotations[0].circles[0].outlineColor.a == 1.0 + assert annots.annotations[0].circles[0].fillColor.r == 0.0 + assert annots.annotations[0].circles[0].fillColor.g == 1.0 + assert annots.annotations[0].circles[0].fillColor.b == 0.0 + assert annots.annotations[0].circles[0].fillColor.a == 1.0 + assert annots.annotations[0].circles[0].thickness == 2 + + +def test_draw_rectangle(): + annotation_helper = AnnotationHelper() + annotation_helper.draw_rectangle( + (0.0, 0.0), (1.0, 1.0), (1.0, 0.0, 0.0, 1.0), (0.0, 1.0, 0.0, 1.0), 2 + ) + annots = annotation_helper.build(timedelta(), 0) + + assert [(i.x, i.y) for i in annots.annotations[0].points[0].points] == [ + (0.0, 0.0), + (1.0, 0.0), + (1.0, 1.0), + (0.0, 1.0), + ] + assert annots.annotations[0].points[0].outlineColor.r == 1.0 + assert annots.annotations[0].points[0].outlineColor.g == 0.0 + assert annots.annotations[0].points[0].outlineColor.b == 0.0 + assert annots.annotations[0].points[0].outlineColor.a == 1.0 + assert annots.annotations[0].points[0].fillColor.r == 0.0 + assert annots.annotations[0].points[0].fillColor.g == 1.0 + assert annots.annotations[0].points[0].fillColor.b == 0.0 + assert annots.annotations[0].points[0].fillColor.a == 1.0 + assert annots.annotations[0].points[0].thickness == 2 + assert annots.annotations[0].points[0].type == dai.PointsAnnotationType.LINE_LOOP + + +def test_draw_text(): + annotation_helper = AnnotationHelper() + annotation_helper.draw_text( + "Test", (0.0, 0.0), (1.0, 0.0, 0.0, 1.0), (0.0, 1.0, 0.0, 1.0), 2 + ) + annots = annotation_helper.build(timedelta(), 0) + + assert annots.annotations[0].texts[0].text == "Test" + assert annots.annotations[0].texts[0].position.x == 0.0 + assert annots.annotations[0].texts[0].position.y == 0.0 + assert annots.annotations[0].texts[0].backgroundColor.r == 0.0 + assert annots.annotations[0].texts[0].backgroundColor.g == 1.0 + assert annots.annotations[0].texts[0].backgroundColor.b == 0.0 + assert annots.annotations[0].texts[0].backgroundColor.a == 1.0 + assert annots.annotations[0].texts[0].textColor.r == 1.0 + assert annots.annotations[0].texts[0].textColor.g == 0.0 + assert annots.annotations[0].texts[0].textColor.b == 0.0 + assert annots.annotations[0].texts[0].textColor.a == 1.0 + assert annots.annotations[0].texts[0].fontSize == 2 + + +def test_draw_points(): + annotation_helper = AnnotationHelper() + annotation_helper.draw_points([(0.0, 0.0), (1.0, 1.0)], (1.0, 0.0, 0.0, 1.0), 2) + annots = annotation_helper.build(timedelta(), 0) + + assert [(i.x, i.y) for i in annots.annotations[0].points[0].points] == [ + (0.0, 0.0), + (1.0, 1.0), + ] + assert annots.annotations[0].points[0].thickness == 2 + assert annots.annotations[0].points[0].type == dai.PointsAnnotationType.POINTS + assert annots.annotations[0].points[0].outlineColor.r == 1.0 + assert annots.annotations[0].points[0].outlineColor.g == 0.0 + assert annots.annotations[0].points[0].outlineColor.b == 0.0 + assert annots.annotations[0].points[0].outlineColor.a == 1.0