Coverage for backpack/annotation/opencv.py: 80%
87 statements
« prev ^ index » next coverage.py v7.2.2, created at 2023-03-30 23:12 +0000
« prev ^ index » next coverage.py v7.2.2, created at 2023-03-30 23:12 +0000
1from typing import Callable, Tuple, Any
3try:
4 import cv2
5except ImportError as e:
6 raise ImportError(
7 'In order to use OpenCVImageAnnotationDriver you should have the optional dependency '
8 'OpenCV installed. You can install it with: pip install "panorama-backpack[opencv]"'
9 )
11import numpy as np
13from ..geometry import Point
14from .driver import AnnotationDriverBase
15from .annotation import (
16 MarkerAnnotation, RectAnnotation, LabelAnnotation, LineAnnotation, PolyLineAnnotation
17)
18from .color import Color
20class OpenCVImageAnnotationDriver(AnnotationDriverBase):
21 ''' Annotation driver implementation for OpenCV images.
23 You should pass an :class:`numpy.ndarray` instance as the context argument of the
24 :meth:`~backpack.annotation.AnnotationDriverBase.render()` method. '''
26 DEFAULT_COLOR = Color(255, 255, 255)
27 DEFAULT_LINEWIDTH = 1
28 DEFAULT_FONT = cv2.FONT_HERSHEY_PLAIN
30 IMG_HEIGHT_FOR_UNIT_FONT_SCALE = 400
31 ''' The height of the image where 1.0 font_scale will be used. '''
33 DEFAULT_TEXT_PADDING = (2, 2)
35 MARKER_STYLE_TO_CV2 = {
36 MarkerAnnotation.Style.CROSS: cv2.MARKER_DIAMOND,
37 MarkerAnnotation.Style.TILTED_CROSS: cv2.MARKER_TILTED_CROSS,
38 MarkerAnnotation.Style.STAR: cv2.MARKER_STAR,
39 MarkerAnnotation.Style.DIAMOND: cv2.MARKER_DIAMOND,
40 MarkerAnnotation.Style.SQUARE: cv2.MARKER_SQUARE,
41 MarkerAnnotation.Style.TRIANGLE_UP: cv2.MARKER_TRIANGLE_UP,
42 MarkerAnnotation.Style.TRIANGLE_DOWN: cv2.MARKER_TRIANGLE_DOWN,
43 }
45 def draw_transparent(self,
46 alpha: float,
47 context: np.ndarray,
48 drawer: Callable[[np.ndarray], None]
49 ) -> None:
50 ''' Semi-transparent drawing.
52 Args:
53 alpha: The transparency factor. Set 0 for opaque drawing, 1 for complete transparency.
54 context: The drawing context.
55 drawer: A callable that will draw an opaque drawing on the passed-in context. This
56 method will make this drawing semi-transparent and render it on the context.
57 '''
58 if alpha == 1.0:
59 drawer(context)
60 elif alpha == 0.0:
61 return
62 else:
63 overlay = context.copy()
64 drawer(overlay)
65 result = cv2.addWeighted(overlay, alpha, context, 1 - alpha, 0)
66 np.copyto(context, result)
68 @staticmethod
69 def scale(point: Any, context: np.ndarray) -> Tuple[float, float]:
70 ''' Converts and scales a point instance to an image context '''
71 x, y = Point.from_value(point)
72 return (int(x * context.shape[1]), int(y * context.shape[0]))
74 @staticmethod
75 def _color_to_cv2(color: Color) -> Tuple[int, int, int]:
76 return (color.b, color.g, color.r)
78 def add_rect(self, anno: RectAnnotation, context: np.ndarray) -> None:
79 color = anno.color or OpenCVImageAnnotationDriver.DEFAULT_COLOR
80 drawer = lambda context: \
81 cv2.rectangle(
82 img=context,
83 pt1=OpenCVImageAnnotationDriver.scale(anno.rect.pt_min, context),
84 pt2=OpenCVImageAnnotationDriver.scale(anno.rect.pt_max, context),
85 color=OpenCVImageAnnotationDriver._color_to_cv2(color),
86 thickness=OpenCVImageAnnotationDriver.DEFAULT_LINEWIDTH
87 )
88 self.draw_transparent(color.alpha, context, drawer)
91 @staticmethod
92 def _get_anchor_shift(
93 horizontal_anchor: LabelAnnotation.HorizontalAnchor,
94 vertical_anchor: LabelAnnotation.VerticalAnchor,
95 size_x: int,
96 size_y: int,
97 baseline: int
98 ) -> Tuple[int, int]:
99 ''' Gets the shift of the text position based on anchor point location and size of the
100 text. '''
101 padding_x, padding_y = OpenCVImageAnnotationDriver.DEFAULT_TEXT_PADDING
102 shift_x, shift_y = (padding_x, -padding_y)
103 if horizontal_anchor == LabelAnnotation.HorizontalAnchor.CENTER:
104 shift_x = -size_x / 2
105 elif horizontal_anchor == LabelAnnotation.HorizontalAnchor.RIGHT:
106 shift_x = -(size_x + padding_x)
107 if vertical_anchor == LabelAnnotation.VerticalAnchor.CENTER:
108 shift_y = size_y / 2
109 elif vertical_anchor == LabelAnnotation.VerticalAnchor.TOP:
110 shift_y = size_y + padding_y
111 elif vertical_anchor == LabelAnnotation.VerticalAnchor.BASELINE:
112 shift_y = baseline
113 return (int(shift_x), int(shift_y))
115 def add_label(self, anno: LabelAnnotation, context: np.ndarray) -> None:
116 ctx_height = context.shape[0]
117 scale = ctx_height / OpenCVImageAnnotationDriver.IMG_HEIGHT_FOR_UNIT_FONT_SCALE * anno.size
118 thickness = max(int(scale), 1)
119 font = OpenCVImageAnnotationDriver.DEFAULT_FONT
120 x, y = OpenCVImageAnnotationDriver.scale(anno.point, context)
121 color = anno.color or OpenCVImageAnnotationDriver.DEFAULT_COLOR
122 if (anno.horizontal_anchor != LabelAnnotation.HorizontalAnchor.LEFT or
123 anno.vertical_anchor != LabelAnnotation.VerticalAnchor.BOTTOM):
124 (size_x, size_y), baseline = cv2.getTextSize(
125 text=anno.text, fontFace=font, fontScale=scale, thickness=thickness
126 )
127 shift_x, shift_y = OpenCVImageAnnotationDriver._get_anchor_shift(
128 anno.horizontal_anchor,
129 anno.vertical_anchor,
130 size_x, size_y, baseline
131 )
132 x += shift_x
133 y += shift_y
135 drawer = lambda context: \
136 cv2.putText(
137 img=context,
138 text=anno.text,
139 org=(x, y),
140 fontFace=font,
141 fontScale=scale,
142 color=OpenCVImageAnnotationDriver._color_to_cv2(color),
143 thickness=thickness
144 )
145 self.draw_transparent(color.alpha, context, drawer)
147 def add_marker(self, anno: MarkerAnnotation, context: np.ndarray) -> None:
148 markerType = OpenCVImageAnnotationDriver.MARKER_STYLE_TO_CV2.get(
149 anno.style, cv2.MARKER_DIAMOND
150 )
151 color = anno.color or OpenCVImageAnnotationDriver.DEFAULT_COLOR
152 drawer = lambda context: \
153 cv2.drawMarker(
154 img=context,
155 position=OpenCVImageAnnotationDriver.scale(anno.point, context),
156 color=OpenCVImageAnnotationDriver._color_to_cv2(color),
157 markerType=markerType
158 )
159 self.draw_transparent(color.alpha, context, drawer)
161 def add_line(self, anno: LineAnnotation, context: Any) -> None:
162 color = anno.color or OpenCVImageAnnotationDriver.DEFAULT_COLOR
163 drawer = lambda context: \
164 cv2.line(
165 img=context,
166 pt1=OpenCVImageAnnotationDriver.scale(anno.line.pt1, context),
167 pt2=OpenCVImageAnnotationDriver.scale(anno.line.pt2, context),
168 color=OpenCVImageAnnotationDriver._color_to_cv2(color),
169 thickness=anno.thickness
170 )
171 self.draw_transparent(color.alpha, context, drawer)
173 def add_polyline(self, anno: PolyLineAnnotation, context: Any) -> None:
174 pts = [OpenCVImageAnnotationDriver.scale(pt, context) for pt in anno.polyline.points]
175 pts = [np.array(pts, dtype=np.int32)]
176 if anno.fill_color is not None:
177 fill_color = OpenCVImageAnnotationDriver._color_to_cv2(anno.fill_color)
178 if anno.polyline.is_convex:
179 drawer = lambda context: \
180 cv2.fillConvexPoly(img=context, points=pts[0], color=fill_color)
181 else:
182 drawer = lambda context: \
183 cv2.fillPoly(img=context, pts=pts, color=fill_color)
184 self.draw_transparent(anno.fill_color.alpha, context, drawer)
186 color = anno.color or OpenCVImageAnnotationDriver.DEFAULT_COLOR
187 drawer = lambda context: \
188 cv2.polylines(
189 img=context,
190 pts=pts,
191 isClosed=anno.polyline.closed,
192 color=OpenCVImageAnnotationDriver._color_to_cv2(color),
193 thickness=anno.thickness
194 )
195 self.draw_transparent(color.alpha, context, drawer)