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

1from typing import Callable, Tuple, Any 

2 

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 ) 

10 

11import numpy as np 

12 

13from ..geometry import Point 

14from .driver import AnnotationDriverBase 

15from .annotation import ( 

16 MarkerAnnotation, RectAnnotation, LabelAnnotation, LineAnnotation, PolyLineAnnotation 

17) 

18from .color import Color 

19 

20class OpenCVImageAnnotationDriver(AnnotationDriverBase): 

21 ''' Annotation driver implementation for OpenCV images. 

22 

23 You should pass an :class:`numpy.ndarray` instance as the context argument of the 

24 :meth:`~backpack.annotation.AnnotationDriverBase.render()` method. ''' 

25 

26 DEFAULT_COLOR = Color(255, 255, 255) 

27 DEFAULT_LINEWIDTH = 1 

28 DEFAULT_FONT = cv2.FONT_HERSHEY_PLAIN 

29 

30 IMG_HEIGHT_FOR_UNIT_FONT_SCALE = 400 

31 ''' The height of the image where 1.0 font_scale will be used. ''' 

32 

33 DEFAULT_TEXT_PADDING = (2, 2) 

34 

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 } 

44 

45 def draw_transparent(self, 

46 alpha: float, 

47 context: np.ndarray, 

48 drawer: Callable[[np.ndarray], None] 

49 ) -> None: 

50 ''' Semi-transparent drawing. 

51 

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) 

67 

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])) 

73 

74 @staticmethod 

75 def _color_to_cv2(color: Color) -> Tuple[int, int, int]: 

76 return (color.b, color.g, color.r) 

77 

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) 

89 

90 

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)) 

114 

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 

134 

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) 

146 

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) 

160 

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) 

172 

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) 

185 

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)