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

1from typing import Any, Optional, Union # noqa F401 

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

3 

4import logging 

5 

6from lunchbox.enforce import Enforce 

7from lunchbox.stopwatch import StopWatch 

8import cv2 

9import numpy as np 

10 

11from cv_depot.core.image import BitDepth, Image 

12from cv_depot.core.color import BasicColor, Color 

13 

14LOGGER = logging.getLogger(__name__) 

15# ------------------------------------------------------------------------------ 

16 

17 

18def gamma(image, value=1.0): 

19 # type: (Image, float) -> Image 

20 ''' 

21 Apply gamma correction to given image. 

22 

23 Args: 

24 image (Image): Image to be modified. 

25 value (int): Gamma value. Default: 1.0. 

26 

27 Raises: 

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

29 EnforceError: If value is less than 0. 

30 

31 Returns: 

32 Image: Gamma adjusted image. 

33 ''' 

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

35 Enforce(value, '>=', 0) 

36 # -------------------------------------------------------------------------- 

37 

38 bit_depth = image.bit_depth 

39 array = image.to_bit_depth(BitDepth.FLOAT32).data 

40 array = array**(1 / value) 

41 output = Image.from_array(array).to_bit_depth(bit_depth) 

42 return output 

43 

44 

45def canny_edges(image, size=0): 

46 # type: (Image, int) -> Image 

47 ''' 

48 Apply a canny edge detection to given image. 

49 

50 Args: 

51 image (Image): Image. 

52 size (int, optional): Amount of dilation applied to result. Default: 0. 

53 

54 Raises: 

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

56 EnforceError: If size is not an integer >= 0. 

57 

58 Returns: 

59 Image: Edge detected image. 

60 ''' 

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

62 Enforce(size, 'instance of', int) 

63 Enforce(size, '>=', 0) 

64 # -------------------------------------------------------------------------- 

65 

66 img = image.to_bit_depth(BitDepth.UINT8).data 

67 img = cv2.Canny(img, 0, 0) 

68 kernel = np.ones((3, 3), np.uint8) 

69 img = cv2.dilate(img, kernel, iterations=size) 

70 output = Image.from_array(img).to_bit_depth(image.bit_depth) 

71 return output 

72 

73 

74def tophat(image, amount, kind='open'): 

75 # type: (Image, int, str) -> Image 

76 ''' 

77 Apply tophat morphology operation to given image. 

78 

79 Args: 

80 image (Image): Image to be modified. 

81 amount (int): Amount of tophat operation. 

82 kind (str, optional): Kind of operation to be performed. 

83 Options include: open, close. Default: open. 

84 

85 Raises: 

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

87 EnforceError: If amount is less than 0. 

88 EnforceError: If kind is not one of: open, close. 

89 

90 Returns: 

91 Image: Image with tophat operation applied. 

92 ''' 

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

94 Enforce(amount, '>=', 0) 

95 msg = 'Illegal tophat kind: {a}. Legal tophat kinds: {b}.' 

96 Enforce(kind, 'in', ['open', 'close'], message=msg) 

97 # -------------------------------------------------------------------------- 

98 

99 lut = dict(open=cv2.MORPH_CLOSE, close=cv2.MORPH_OPEN) 

100 opt = lut[kind] 

101 bit_depth = image.bit_depth 

102 img = image.to_bit_depth(BitDepth.FLOAT32).data 

103 kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (3, 3)) 

104 arr = cv2 \ 

105 .morphologyEx(img, opt, kernel, iterations=amount) \ 

106 .astype(np.float32) 

107 output = Image.from_array(arr).to_bit_depth(bit_depth) 

108 return output 

109 

110 

111def linear_lookup(lower=0, upper=1): 

112 # type: (float, float) -> np.vectorize 

113 r''' 

114 Generates a linear lookup table with an upper and lower shoulder. 

115 

116 .. image:: images/linear_lut.png 

117 

118 Args: 

119 lower (float, optional): Lower shoulder value. Default: 0. 

120 upper (float, optional): Upper shoulder value. Default: 1. 

121 

122 Returns: 

123 numpy.vectorize: Anonymous function that applies lut elementwise to a 

124 given numpy array. 

125 ''' 

126 lut = lambda x: min(max((x - lower), 0) * (1 / (upper - lower)), 1) 

127 return np.vectorize(lut) 

128 

129 

130def linear_smooth(image, blur=3, lower=0, upper=1): 

131 # type: (Image, int, float, float) -> Image 

132 ''' 

133 Blur given image then apply linear thresholding it. 

134 

135 Args: 

136 image (Image): Image matte to be smoothed. 

137 blur (int, optional): Size of blur. Default: 3. 

138 lower (float, optional): Lower shoulder value. Default: 0. 

139 upper (float, optional): Upper shoulder value. Default: 1. 

140 

141 Raises: 

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

143 EnforceError: If blur is less than 0. 

144 EnforceError: If lower or upper is less than 0. 

145 EnforceError: If lower or upper is greater than 1. 

146 EnforceError: If lower is greater than upper. 

147 

148 Returns: 

149 Image: Smoothed image. 

150 ''' 

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

152 Enforce(blur, '>=', 0) 

153 Enforce(lower, '>=', 0) 

154 Enforce(lower, '<=', 1) 

155 Enforce(upper, '>=', 0) 

156 Enforce(upper, '<=', 1) 

157 Enforce(lower, '>=', 0) 

158 msg = 'Lower bound cannot be greater than upper bound. {a} > {b}' 

159 Enforce(lower, '<=', upper, message=msg) 

160 # -------------------------------------------------------------------------- 

161 

162 bit_depth = image.bit_depth 

163 img = image.to_bit_depth(BitDepth.FLOAT32).data 

164 img = cv2.blur(img, (blur, blur)) 

165 lut = linear_lookup(lower=lower, upper=upper) 

166 img = lut(img).astype(np.float32) 

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

168 return output 

169 

170 

171def key_exact_color(image, color, channel='a', invert=False): 

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

173 ''' 

174 Keys given image according to the color of its pixels values. 

175 Where that pixel color exactly matches the given color, the mask channel 

176 will be 1, otherwise it will be 0. 

177 

178 Args: 

179 image (Image): Image to be evaluated. 

180 color (Color or BasicColor): Color to be used for masking. 

181 channel (str, optional): Mask channel name. Default: a. 

182 invert (bool, optional): Whether to invert the mask. Default: False. 

183 

184 Raises: 

185 EnforceError: If image is not an Image instance. 

186 EnforceError: If channel is not a string. 

187 EnforceError: If invert is not a boolean. 

188 EnforceError: If RGB is not found in image channels. 

189 

190 Returns: 

191 Image: Image with mask channel. 

192 ''' 

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

194 Enforce(channel, 'instance of', str) 

195 Enforce(invert, 'instance of', bool) 

196 # -------------------------------------------------------------------------- 

197 

198 # get color 

199 if isinstance(color, str): 

200 color = BasicColor.from_string(color) 

201 if isinstance(color, BasicColor): 

202 color = Color.from_basic_color(color) 

203 clr = color.to_array() 

204 

205 # determine num channels 

206 rgb = list('rgb') 

207 img = image.to_bit_depth(BitDepth.FLOAT32) 

208 if image.num_channels == 1: 

209 x = img.data[..., np.newaxis] 

210 x = np.concatenate([x, x, x], axis=2) 

211 img = Image.from_array(x) 

212 else: 

213 diff = sorted(list(set(rgb).difference(image.channels))) 

214 msg = f'{diff} not found in image channels. ' 

215 msg += f'Given channels: {image.channels}.' 

216 Enforce(len(diff), '==', 0, message=msg) 

217 

218 # create mask 

219 mask = np.equal(clr, img[:, :, rgb].data) 

220 mask = np.apply_along_axis(all, 2, mask) \ 

221 .astype(np.float32)[..., np.newaxis] 

222 if invert: 

223 mask = -1 * mask + 1 

224 

225 # add mask to image 

226 chans = image.channels 

227 if image.num_channels == 1: 

228 arr = img[:, :, 'r'].data[..., np.newaxis] 

229 else: 

230 chans = list(filter(lambda x: x != channel, img.channels)) 

231 arr = img[:, :, chans].data 

232 

233 arr = np.concatenate([arr, mask], axis=2) 

234 output = Image.from_array(arr).to_bit_depth(image.bit_depth) 

235 output = output.set_channels(chans + [channel]) 

236 return output 

237 

238 

239def kmeans( 

240 image, # type: Image 

241 num_centroids=10, # type: int 

242 centroids=None, # type: Optional[list[tuple[int, int, int]]] 

243 max_iterations=100, # type: int 

244 accuracy=1.0, # type: float 

245 epochs=10, # type: int 

246 seeding='random', # type: str 

247 generate_report=False # type: bool 

248): # type: (...) -> Union[Image, tuple[Image, dict]] 

249 ''' 

250 Applies k-means to the given image. 

251 

252 Args: 

253 image (Image): Image instance. 

254 num_centroids (int, optional): Number of centroids to use. Default: 10. 

255 centroids (list, optional): List of triplets. Default: None. 

256 max_iterations (int, optional): Maximum number of k-means updates 

257 allowed per centroid. Default: 100. 

258 accuracy (float, optional): Minimum accuracy required of clusters. 

259 epochs (int, optional): Number of times algorithm is applied with 

260 different initial labels. Default: 10. 

261 seeding (str, optional): How intial centroids are generated. Default: 

262 random. Options include: random, pp_centers. 

263 generate_report (bool, optional): If true returns report in addition to 

264 image. Default: False. 

265 

266 Raises: 

267 EnforceError: If image is not an Image instance. 

268 ValueError: If invalid seeding option is given. 

269 

270 Returns: 

271 Image or tuple[Image, dict]: K-means image or K-means image and K-means 

272 report. 

273 ''' 

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

275 

276 stopwatch = StopWatch() 

277 stopwatch.start() 

278 # -------------------------------------------------------------------------- 

279 

280 source_bit_depth = image.bit_depth 

281 data = image.to_bit_depth(BitDepth.FLOAT32)\ 

282 .data.reshape((-1, image.num_channels)) 

283 

284 seed = None 

285 if seeding == 'random': 

286 seed = cv2.KMEANS_RANDOM_CENTERS 

287 elif seeding == 'pp_centers': 

288 seed = cv2.KMEANS_PP_CENTERS 

289 else: 

290 msg = f'{seeding} is not a valid seeding option. Options include: ' 

291 msg += '[random, pp_centers].' 

292 raise ValueError(msg) 

293 

294 # terminate centroid updates when max iteration or min accuracy is achieved 

295 crit = cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER 

296 

297 if centroids is not None: 

298 num_centroids = len(centroids) 

299 

300 compactness, labels, centroids = cv2.kmeans( 

301 data=data, 

302 K=num_centroids, 

303 bestLabels=centroids, 

304 criteria=(crit, max_iterations, accuracy), 

305 attempts=epochs, 

306 flags=seed, 

307 ) # type: ignore 

308 

309 centroids_ = centroids # type: Any 

310 output = np.float32(centroids_)[labels.flatten()] # type: ignore 

311 output = output.reshape(image.data.shape) 

312 output = Image.from_array(output).to_bit_depth(source_bit_depth) 

313 

314 if generate_report: 

315 report = dict( 

316 compactness=compactness, 

317 labels=labels, 

318 centroids=centroids 

319 ) 

320 stopwatch.stop() 

321 LOGGER.warning(f'KmeansRuntime: {stopwatch.human_readable_delta}.') 

322 return output, report 

323 

324 stopwatch.stop() 

325 LOGGER.warning(f'Kmeans Runtime: {stopwatch.human_readable_delta}.') 

326 return output