Coverage for /home/ubuntu/cv-depot/python/cv_depot/ops/draw.py: 99%

1from typing import Any, List, Tuple, Union # noqa F401 

2from cv_depot.core.types import AnyColor # noqa F401 

3 

4import copy 

5 

6from lunchbox.enforce import Enforce 

7import cv2 

8import numpy as np 

9 

10from cv_depot.core.channel_map import ChannelMap 

11from cv_depot.core.color import BasicColor, Color 

12from cv_depot.core.image import BitDepth, Image 

13import cv_depot.ops.channel as cvchan 

14import cv_depot.ops.draw as cvdraw 

15import cv_depot.ops.filter as cvfilt 

16# ------------------------------------------------------------------------------ 

17 

18 

19def swatch( 

20 shape, # type: Union[Tuple[int, int, int], List[int]] 

21 color, # type: AnyColor 

22 fill_value=0.0, # type: float 

23 bit_depth=BitDepth.FLOAT32 # type: BitDepth 

24): 

25 # type: (...) -> Image 

26 ''' 

27 Creates an image of the given shape and color. 

28 

29 Args: 

30 shape (tuple[int]): List of 3 ints. 

31 color (Color): Color of swatch RGB or L channels. 

32 fill_value (float, optional): Value to fill additional channels with. 

33 Default: 0. 

34 bit_depth (BitDepth): Bit depth of swatch. Default: BitDepth.FLOAT32 

35 

36 Raises: 

37 EnforceError: If shape is not a tuple of 3 integers. 

38 EnforceError: If shape has any zero dimensions. 

39 

40 Returns: 

41 Image: Color swatch. 

42 ''' 

43 msg = f'Shape must be an tuple or list of 3 integers. Given shape: {shape}. ' 

44 msg += f'Given type: {type(shape)}.' 

45 Enforce(shape.__class__.__name__, 'in', ['tuple', 'list'], message=msg) 

46 Enforce(len(shape), '==', 3, message=msg) 

47 for item in shape: 

48 Enforce(item, 'instance of', int, message=msg) 

49 

50 msg = 'Illegal shape: Each shape dimension must be greater than 0. ' 

51 msg += f'Given shape: {shape}.' 

52 Enforce(min(shape), '>', 0, message=msg) 

53 # -------------------------------------------------------------------------- 

54 

55 w, h, c = shape 

56 

57 if isinstance(color, str): 

58 color = BasicColor.from_string(color) 

59 if isinstance(color, BasicColor): 

60 color = Color.from_basic_color(color) 

61 color = Color.from_array( 

62 color.to_array(), num_channels=c, fill_value=fill_value 

63 ) 

64 

65 output = np.ones((h, w, c), np.float32) * color.to_array() 

66 output = Image.from_array(output).to_bit_depth(bit_depth) 

67 return output 

68 

69 

70def grid(image, shape, color=BasicColor.CYAN.name, thickness=1): 

71 # type: (Image, Tuple[int, int], AnyColor, int) -> Image 

72 ''' 

73 Draws a grid on a given image. 

74 

75 Args: 

76 image (Image): Image to be drawn on. 

77 shape (tuple[int]): Width, height tuple. 

78 color (Color, optional): Color of grid. Default: BasicColor.CYAN. 

79 thickness (int, optional): Thickness of grid lines in pixels. 

80 Default: 1. 

81 

82 Raises: 

83 EnforceError: If image is not an Image instance. 

84 EnforceError: If shape is not of length 2. 

85 EnforceError: If width or height of shape is less than 0: 

86 EnforceError: If color is not an instance of Color or BasicColor. 

87 EnforceError: If thickness is not greater than 0. 

88 

89 Returns: 

90 Image: Image with grid. 

91 ''' 

92 msg = 'Illegal image. {a.__class__.__name__} is not an instance of Image.' 

93 Enforce(image, 'instance of', Image, message=msg) 

94 msg = f'Illegal shape. Expected (w, h). Found: {shape}.' 

95 Enforce(len(shape), '==', 2, message=msg) 

96 msg = 'Shape width must be greater than or equal to 0. {a} < 0.' 

97 Enforce(shape[0], '>=', 0, message=msg) 

98 msg = 'Shape height must be greater than or equal to 0. {a} < 0.' 

99 Enforce(shape[1], '>=', 0, message=msg) 

100 msg = 'Color type must be in {b}. Found: {a}.' 

101 Enforce(color.__class__.__name__, 'in', ['Color', 'BasicColor'], message=msg) 

102 msg = 'Line thickness must be an integer. Found: {a.__class__.__name__}.' 

103 Enforce(thickness, 'instance of', int, message=msg) 

104 msg = 'Line thickness must be greater than 0. {a} !> 0.' 

105 Enforce(thickness, '>', 0) 

106 # -------------------------------------------------------------------------- 

107 

108 w0, h0 = shape 

109 w, h = image.width_and_height 

110 bit_depth = image.bit_depth 

111 img = image.to_bit_depth(BitDepth.FLOAT32).data.copy() 

112 if isinstance(color, str): 

113 color = BasicColor.from_string(color) 

114 if isinstance(color, BasicColor): 

115 color = Color.from_basic_color(color) 

116 clr = color.to_array().tolist() # type: Any 

117 

118 # vertical lines 

119 w_step = int(round(w / w0, 0)) 

120 for x in range(w_step, w, w_step): 

121 img = cv2.line(img, (x, 0), (x, h), clr, thickness) 

122 

123 # horizontal lines 

124 h_step = int(round(h / h0, 0)) 

125 for y in range(h_step, h, h_step): 

126 img = cv2.line(img, (0, y), (w, y), clr, thickness) 

127 

128 output = Image.from_array(img).to_bit_depth(bit_depth) 

129 return output 

130 

131 

132def checkerboard(tiles_wide, tiles_high, tile_shape=(10, 10)): 

133 # type: (int, int, Tuple[int, int]) -> Image 

134 ''' 

135 Draws a checkerboard of given width, height and tile shape. 

136 

137 Args: 

138 tiles_wide (int): Number of tiles wide checkerboard will be. 

139 tiles_high (int): Number of tiles high checkerboard will be. 

140 tile_shape (tuple[int], optional): Width, height tuple of tile shape. 

141 Default: (10, 10). 

142 

143 Raises: 

144 EnforceError: If tiles_wide is not greater than 0. 

145 EnforceError: If tiles_high is not greater than 0. 

146 EnforceError: If tile width is not greater than 0. 

147 EnforceError: If tile height is not greater than 0. 

148 

149 Returns: 

150 Image: Checkerboard image. 

151 ''' 

152 w, h = tile_shape 

153 shape = (w, h, 3) 

154 

155 msg = 'Tiles_wide must be greater than 0. {a} !> 0.' 

156 Enforce(tiles_wide, '>', 0, message=msg) 

157 msg = 'Tiles_high must be greater than 0. {a} !> 0.' 

158 Enforce(tiles_high, '>', 0, message=msg) 

159 msg = 'Tile width must be greater than 0. {a} !> 0.' 

160 Enforce(w, '>', 0, message=msg) 

161 msg = 'Tile height must be greater than 0. {a} !> 0.' 

162 Enforce(h, '>', 0, message=msg) 

163 # -------------------------------------------------------------------------- 

164 

165 black = swatch(shape, BasicColor.BLACK).data 

166 white = swatch(shape, BasicColor.WHITE).data 

167 even = [] 

168 odd = [] 

169 for i in range(0, tiles_wide): 

170 if i == 0 or i % 2 == 0: 

171 even.append(black) 

172 odd.append(white) 

173 else: 

174 even.append(white) 

175 odd.append(black) 

176 

177 even = np.concatenate(even, axis=1) 

178 odd = np.concatenate(odd, axis=1) 

179 

180 rows = [] 

181 for i in range(0, tiles_high): 

182 if i == 0 or i % 2 == 0: 

183 rows.append(even) 

184 else: 

185 rows.append(odd) 

186 

187 output = Image.from_array(np.concatenate(rows, axis=0)) 

188 return output 

189 

190 

191def highlight( 

192 image, mask='a', opacity=0.5, color=BasicColor.CYAN2.name, inverse=False 

193): 

194 # type: (Image, str, float, AnyColor, bool) -> Image 

195 ''' 

196 Highlight a masked portion of a given image according to a given channel. 

197 

198 Args: 

199 image (Image): Image to be highlighted. 

200 mask (str, optional): Channel to be used as mask. Default: alpha. 

201 opacity (float, optional): Opacity of highlight overlayed on image. 

202 Default: 0.5 

203 color (Color or BasicColor, optional): Color of highlight. 

204 Default: BasicColor.CYAN2. 

205 inverse (bool, optional): Whether to invert the highlight. 

206 Default: False. 

207 

208 Raises: 

209 EnforceError: If image is not an instance of Image. 

210 EnforceError: If mask is not an instance of str. 

211 EnforceError: If mask not found in image channels. 

212 EnforceError: If opacity is < 0 or > 1. 

213 EnforceError: If color is not instance of Color. 

214 EnforceError: If inverse is not a boolean. 

215 

216 Returns: 

217 Image: Highlighted image. 

218 ''' 

219 if isinstance(color, str): 

220 color = BasicColor.from_string(color) 

221 if isinstance(color, BasicColor): 

222 color = Color.from_basic_color(color) 

223 

224 Enforce(image, 'instance of', Image) 

225 Enforce(mask, 'instance of', str) 

226 msg = 'Mask channel: {a} not found in image channels: {b}.' 

227 Enforce(mask, 'in', image.channels, message=msg) 

228 Enforce(opacity, '>=', 0) 

229 Enforce(opacity, '<=', 1) 

230 Enforce(color, 'instance of', Color) 

231 Enforce(inverse, 'instance of', bool) 

232 # -------------------------------------------------------------------------- 

233 

234 img = image.to_bit_depth(BitDepth.FLOAT32) 

235 channels = copy.deepcopy(image.channels) 

236 matte = cvchan.remap_single_channel(img[:, :, mask], channels) 

237 imatte = cvchan.invert(matte) 

238 if inverse: 

239 matte, imatte = imatte, matte 

240 

241 swatch = cvdraw.swatch(image.shape, color, fill_value=1.0) 

242 data = (image.data * imatte.data) + (swatch.data * matte.data) 

243 output = Image.from_array(data).to_bit_depth(image.bit_depth) 

244 output = cvchan \ 

245 .mix(image, output, amount=1 - opacity) \ 

246 .set_channels(image.channels) 

247 return output 

248 

249 

250def outline(image, mask='a', width=10, color=BasicColor.CYAN2.name): 

251 # type (Image, str, int, AnyColor) -> Image 

252 ''' 

253 Use a given mask to outline a given image. 

254 

255 Args: 

256 image (Image): Image with mask channel. 

257 mask (str, optional): Mask channel. Default: alpha. 

258 width (int, optional): Outline width. Default: 10. 

259 color (Color or BasicColor, optional): Color of outline. 

260 Default: BasicColor.CYAN2 

261 

262 Raises: 

263 EnforceError: If image is not an instance of Image. 

264 EnforceError: If channel not found in image channels. 

265 EnforceError: If width is not >= 0. 

266 EnforceError: If color is not an instance of Color or BasicColor. 

267 

268 Returns: 

269 Image: Image with outline. 

270 ''' 

271 Enforce(image, 'instance of', Image) 

272 msg = 'Mask channel: {a} not found in image channels: {b}.' 

273 Enforce(mask, 'in', image.channels, message=msg) 

274 Enforce(width, '>=', 0) 

275 # -------------------------------------------------------------------------- 

276 

277 cmap = ChannelMap({c: f'0.{c}' for c in image.channels}) 

278 cmap[mask] = '1.l' 

279 w = int(round(width / 2, 0)) 

280 edge = cvfilt.canny_edges(image[:, :, mask], size=w) 

281 output = cvchan.remap([image, edge], cmap) 

282 output = highlight(output, mask=mask, opacity=1.0, color=color) 

283 return output 

284 

285 

286def annotate( 

287 image, # type: Image 

288 mask='a', # type: str 

289 opacity=0.5, # type: float 

290 width=10, # type: int 

291 color=BasicColor.CYAN2.name, # type: AnyColor 

292 inverse=False, # type: bool 

293 keep_mask=False, # type: bool 

294): 

295 # type (...) -> Image 

296 ''' 

297 Annotate a given image according to a given mask channel. 

298 

299 Args: 

300 image (Image): Image with mask channel. 

301 mask (str, optional): Mask channel. Default: alpha. 

302 opacity (float, optional): Opacity of annotation. Default: 0.5 

303 width (int, optional): Outline width. Default: 10. 

304 color (Color or BasicColor, optional): Color of outline. 

305 Default: BasicColor.CYAN2 

306 inverse (bool, optional): Whether to invert the annotation. 

307 Default: False. 

308 keep_mask (bool, optional): Whether to keep the mask channel. 

309 Default: False. 

310 

311 Returns: 

312 Image: Image with outline. 

313 ''' 

314 output = highlight( 

315 image, mask=mask, opacity=opacity, color=color, inverse=inverse 

316 ) 

317 output = outline(output, mask=mask, width=width, color=color) 

318 if not keep_mask: 

319 channels = copy.deepcopy(image.channels) 

320 channels.remove(mask) 

321 output = output[:, :, channels] 

322 return output