Coverage for /home/ubuntu/cv-depot/python/cv_depot/ops/channel.py: 100%

1from typing import Union # noqa F401 

2 

3import re 

4 

5from lunchbox.enforce import Enforce 

6import cv2 

7import numpy as np 

8 

9from cv_depot.core.channel_map import ChannelMap 

10from cv_depot.core.color import BasicColor 

11from cv_depot.core.enum import BitDepth 

12from cv_depot.core.image import Image 

13import cv_depot.core.image as cvimg 

14import cv_depot.ops.draw as cvdraw 

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

16 

17 

18def has_super_brights(image): 

19 # type: (Image) -> bool 

20 ''' 

21 Determines if given image has values above 1.0. 

22 

23 Args: 

24 image (Image): Image instance. 

25 

26 Raises: 

27 EnforceError: If image is not an Image instance. 

28 

29 Returns: 

30 bool: Presence of super brights. 

31 ''' 

32 return cvimg._has_super_brights(image) 

33 

34 

35def has_super_darks(image): 

36 # type: (Image) -> bool 

37 ''' 

38 Determines if given image has values below 0.0. 

39 

40 Args: 

41 image (Image): Image instance. 

42 

43 Raises: 

44 EnforceError: If image is not an Image instance. 

45 

46 Returns: 

47 bool: Presence of super darks. 

48 ''' 

49 return cvimg._has_super_darks(image) 

50 

51 

52def to_hsv(image): 

53 # type: (Image) -> Image 

54 ''' 

55 Convert image to hue, saturation, value colorspace. 

56 

57 Args: 

58 Image: Image to be converted. 

59 

60 Raises: 

61 AttributeError: If given image does not have RGB channels. 

62 

63 Returns: 

64 Image: Image converted to HSV. 

65 ''' 

66 rgb = list('rgb') 

67 channels = set(image.channels).intersection(rgb) 

68 if channels != set(rgb): 

69 msg = 'Image does not contain RGB channels. ' 

70 msg += f'Channels found: {image.channels}.' 

71 raise AttributeError(msg) 

72 

73 img_ = image.to_bit_depth(BitDepth.FLOAT32) 

74 img = img_[:, :, rgb].data 

75 img = cv2.cvtColor(img, cv2.COLOR_RGB2HSV) 

76 img[:, :, 0] = np.divide(img[:, :, 0], 360) 

77 output = Image\ 

78 .from_array(img)\ 

79 .set_channels(list('hsv'))\ 

80 .to_bit_depth(image.bit_depth) 

81 return output 

82 

83 

84def to_rgb(image): 

85 # type: (Image) -> Image 

86 ''' 

87 Convert image from HSV to RGB. 

88 

89 Args: 

90 Image: Image to be converted. 

91 

92 Raises: 

93 AttributeError: If given image does not have RGB channels. 

94 

95 Returns: 

96 Image: Image converted to RGB. 

97 ''' 

98 hsv = list('hsv') 

99 channels = set(image.channels).intersection(hsv) 

100 if channels != set(hsv): 

101 msg = 'Image does not contain HSV channels. ' 

102 msg += f'Channels found: {image.channels}.' 

103 raise AttributeError(msg) 

104 

105 img_ = image.to_bit_depth(BitDepth.FLOAT32) 

106 img = img_[:, :, hsv].data 

107 img = cv2.cvtColor(img, cv2.COLOR_HSV2RGB) 

108 img[:, :, 0] = np.divide(img[:, :, 0], 360) 

109 output = Image\ 

110 .from_array(img)\ 

111 .set_channels(list('rgb'))\ 

112 .to_bit_depth(image.bit_depth) 

113 return output 

114 

115 

116def invert(image): 

117 # type: (Image) -> Image 

118 ''' 

119 Inverts the values of the given image. 

120 Black becomes white, white becomes black. 

121 

122 Args: 

123 image (Image): Image to be inverted. 

124 

125 Raises: 

126 EnforeError: If image is not an instance of Image. 

127 

128 Returns: 

129 image: Image 

130 ''' 

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

132 # -------------------------------------------------------------------------- 

133 

134 bit_depth = image.bit_depth 

135 data = image.to_bit_depth(BitDepth.FLOAT32).data 

136 data = data * -1 + 1 

137 output = Image.from_array(data).to_bit_depth(bit_depth) 

138 return output 

139 

140 

141def mix(a, b, amount=0.5): 

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

143 ''' 

144 Mix images A and B by a given amount. 

145 An amount of 1.0 means 100% of image A. 

146 An amount of 0.0 means 100% of image B. 

147 

148 Args: 

149 a (Image): Image A. 

150 b (Image): Image B. 

151 amount (float, optional): Amount of image A. Default: 0.5 

152 

153 Raises: 

154 EnforceError: If a is not an Image instance. 

155 EnforceError: If b is not an Image instance. 

156 EnforceError: If amount is not between 0 and 1. 

157 

158 Returns: 

159 Image: Mixture of A and B. 

160 ''' 

161 Enforce(a, 'instance of', Image) 

162 Enforce(b, 'instance of', Image) 

163 Enforce(amount, '<=', 1) 

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

165 # -------------------------------------------------------------------------- 

166 

167 amount = float(amount) 

168 bit_depth = a.bit_depth 

169 x = a.to_bit_depth(BitDepth.FLOAT32).data 

170 y = b.to_bit_depth(BitDepth.FLOAT32).data 

171 img = x * amount + y * (1 - amount) 

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

173 return output 

174 

175 

176def remap_single_channel(image, channels): 

177 # type: (Image, list) -> Image 

178 ''' 

179 Maps an image with a single channel to an image of a given number of 

180 channels. 

181 

182 Args: 

183 image (Image): Image to be mapped. 

184 channels (list): List of channel names to map image to. 

185 

186 Raises: 

187 EnforceError: If image is not an Image with only one channel. 

188 EnforceError: If channels is not a list. 

189 

190 Returns: 

191 Image: Image with given channels. 

192 ''' 

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

194 msg = 'Image must be an Image with only 1 channel. ' 

195 msg += 'Channels found: {a}.' 

196 Enforce(image.num_channels, '==', 1, message=msg) 

197 Enforce(channels, 'instance of', list) 

198 # -------------------------------------------------------------------------- 

199 

200 data = np.squeeze(image.data)[..., np.newaxis] 

201 output = np.concatenate([data] * len(channels), axis=2) 

202 output = Image.from_array(output).set_channels(channels) 

203 return output 

204 

205 

206def remap(images, channel_map): 

207 # type: (Union[Image, list[Image]], ChannelMap) -> Image 

208 ''' 

209 Maps many images into a single image according to a given channel 

210 map. 

211 

212 Args: 

213 images (Image, list[Image]): Images. 

214 channel_map (ChannelMap): Mapping of image channels into output image. 

215 

216 Raises: 

217 EnforceError: If images is not an instance of Image or list of Images. 

218 EnforceError: If images are not of all the same width and height. 

219 EnforceError: If channel_map is not an instance of ChannelMap. 

220 

221 Returns: 

222 Image: Combined image. 

223 ''' 

224 if isinstance(images, Image): 

225 images = [images] 

226 

227 msg = 'Images must be an Image or list of Images of uniform width and height.' 

228 msg += f' Given type: {images.__class__.__name__}' 

229 [Enforce(x, 'instance of', Image) for x in images] 

230 

231 shapes = {x.width_and_height for x in images} 

232 Enforce(len(shapes), '==', 1, message=msg) 

233 

234 Enforce(channel_map, 'instance of', ChannelMap) 

235 # -------------------------------------------------------------------- 

236 

237 w, h = images[0].width_and_height 

238 bit_depth = images[0].bit_depth 

239 images = [x.to_bit_depth(bit_depth) for x in images] 

240 

241 shape = (w, h, 1) 

242 black = cvdraw.swatch(shape, BasicColor.BLACK, bit_depth=bit_depth) 

243 white = cvdraw.swatch(shape, BasicColor.WHITE, bit_depth=bit_depth) 

244 

245 channels = [] 

246 for chan in channel_map.source: 

247 chan_l = chan.lower() 

248 img = None 

249 

250 if chan_l in ['b', 'black']: 

251 img = black 

252 elif chan_l in ['w', 'white']: 

253 img = white 

254 else: 

255 frame, tgt_chan = re.split(r'\.', chan, maxsplit=1) 

256 img = images[int(frame)][:, :, tgt_chan] 

257 

258 array = np.squeeze(img.data)[..., np.newaxis] 

259 channels.append(array) 

260 

261 img = np.concatenate(channels, axis=2).astype(bit_depth.dtype) 

262 output = Image.from_array(img).set_channels(channel_map.target) # type: ignore 

263 return output