Coverage for python / lsst / pipe / tasks / prettyPictureMaker / _colorMapper.py: 9%

158 statements  

« prev     ^ index     » next       coverage.py v7.12.0, created at 2025-12-20 01:16 -0800

1__all__ = ("mapUpperBounds", "latLum", "colorConstantSat", "lsstRGB", "mapUpperBounds") 

2 

3import logging 

4import numpy as np 

5import skimage 

6import colour 

7import cv2 

8from skimage.restoration import inpaint_biharmonic 

9 

10from ._localContrast import localContrast, makeGaussianPyramid, makeLapPyramid, levelPadder 

11from lsst.cpputils import fixGamutOK 

12 

13from numpy.typing import NDArray 

14from typing import Callable, Mapping, Literal 

15 

16 

17def latLum( 

18 values, 

19 stretch: float = 400, 

20 max: float = 1, 

21 floor: float = 0.00, 

22 Q: float = 0.7, 

23 doDenoise: bool = False, 

24 highlight: float = 1.0, 

25 shadow: float = 0.0, 

26 midtone: float = 0.5, 

27) -> NDArray: 

28 """ 

29 Scale the input luminosity values to maximize the dynamic range visible. 

30 

31 Parameters 

32 ---------- 

33 values : `NDArray` 

34 The input image luminosity data of of. 

35 stretch : `float`, optional 

36 A parameter for the arcsinh function. 

37 max : `float`, optional 

38 Maximum value for intensity scaling on a scale of 0-1. 

39 floor : `float`, optional 

40 A value added to each pixel in arcsinh transform, this ensures values in 

41 the arcsinh transform is no smaller than the supplied value. 

42 Q : `float`, optional 

43 Another parameter for the arcsinh function and scaling factor for 

44 softening. 

45 doDenoise : `bool`, optional 

46 Denoise the image if desired. 

47 highlight : `float` 

48 This is the value (between 0 and 1) that maps to be "white" in the 

49 output. Decreasing this makes fainter things more luminous but 

50 clips the brightest end. This is a linear transform applied to the 

51 values after arcsinh. 

52 shadow : `float` 

53 This is the value (between 0 and 1) that maps to be "black" in the 

54 output. Increasing this makes fainter things darker but 

55 clips the lowest values. This is a linear transform applied to the 

56 values after arcsinh. 

57 midtone : `float` 

58 This is the value (between 0 and 1) that adjusts the balance between 

59 white and black. Decreasing this makes fainter things more luminous, 

60 increasing does the opposite. This is a linear transform applied to the 

61 values after arcsinh. 

62 

63 Returns: 

64 luminance : `NDArray` 

65 The stretched luminosity data. 

66 """ 

67 

68 # De-noise the input image using wavelet de-noising. 

69 if doDenoise: 

70 values = skimage.restoration.denoise_wavelet(values) 

71 values = abs(values) 

72 

73 # Scale values from 0-1 to 0-100 as various algorithm expect that range. 

74 values *= 100 

75 

76 # Calculate the slope for arcsinh transformation based on Q and stretch 

77 # parameters. 

78 slope = 0.1 * 100 / np.arcsinh(0.1 * Q) 

79 

80 # Apply the modified luminosity transformation using arcsinh function. 

81 soften = Q / stretch 

82 intensities = np.arcsinh((abs(values) * soften + floor) * slope) 

83 

84 # Always normalize by what the original max value (100) scales to. 

85 maximum_intensity = np.arcsinh((100 * soften + floor) * slope) 

86 

87 intensities /= maximum_intensity 

88 

89 # Scale the intensities with linear manipulation for contrast 

90 intensities = (intensities - shadow) / (highlight - shadow) 

91 intensities = ((midtone - 1) * intensities) / (((2 * midtone - 1) * intensities) - midtone) 

92 

93 np.clip(intensities, 0, max, out=intensities) 

94 

95 # Reset the input array. 

96 values /= 100 

97 

98 return intensities 

99 

100 

101def mapUpperBounds(img: NDArray, quant: float = 0.9, absMax: float | None = None) -> NDArray: 

102 """Bound images to a range between zero and one. 

103 

104 Some images supplied aren't properly bounded with a maximum value of 1. 

105 Either the images exceed the bounds of 1, or that no value seems to close, 

106 implying indeterminate maximum value. This function determines an 

107 appropriate maximum either by taking the value supplied in the absMax 

108 argument or by scaling the maximum across all channels with the 

109 supplied quant variable. 

110 

111 Parameters 

112 ---------- 

113 img : `NDArray` like 

114 Must have dimensions of y,x,3 where the channels are in RGB order 

115 quant : `float` 

116 Value to scale the maximum pixel value, in any channel, by to 

117 determine the maximum flux allowable in all channels. Ignored 

118 if absMax isn't None. 

119 absMax : `float` or `None` 

120 If this value is not None, use it as the maximum pixel value 

121 for all channels, unless scaleBoundFactor is set in which case 

122 it is only the maximum if the value determined from the image 

123 and quant is larger than scaleBoundFactor times absMax. This is 

124 to prevent individual frames in a mosaic from being scaled too 

125 faint if absMax is too large for one region. 

126 

127 Returns 

128 ------- 

129 image : `NDArray` 

130 The result of the remapping process 

131 """ 

132 if np.max(img) == 1: 

133 return img 

134 

135 r = img[:, :, 0] 

136 g = img[:, :, 1] 

137 b = img[:, :, 2] 

138 

139 r_quant = np.quantile(r, 0.95) 

140 g_quant = np.quantile(g, 0.95) 

141 b_quant = np.quantile(b, 0.95) 

142 turnover = np.max((r_quant, g_quant, b_quant)) 

143 

144 if absMax is not None: 

145 scale = absMax 

146 else: 

147 scale = turnover * quant 

148 

149 image = np.empty(img.shape) 

150 image[:, :, 0] = r / scale 

151 image[:, :, 1] = g / scale 

152 image[:, :, 2] = b / scale 

153 

154 # Clip values that exceed the bound to ensure all values are within [0, absMax] 

155 np.clip(image, 0, 1, out=image) 

156 

157 return image 

158 

159 

160def colorConstantSat( 

161 oldLum: NDArray, 

162 luminance: NDArray, 

163 a: NDArray, 

164 b: NDArray, 

165 saturation: float = 0.6, 

166 maxChroma: float = 80, 

167) -> tuple[NDArray, NDArray]: 

168 """ 

169 Adjusts the color saturation while keeping the hue constant. 

170 

171 This function adjusts the chromaticity (a, b) of colors to maintain a 

172 consistent saturation level, based on their original luminance. It uses 

173 the CIELAB color space representation and the `luminance` is the new target 

174 luminance for all colors. 

175 

176 Parameters 

177 ---------- 

178 oldLum : `NDArray` 

179 Luminance values of the original colors. 

180 luminance : `NDArray` 

181 Target luminance values for the transformed colors. 

182 a : `NDArray` 

183 Chromaticity parameter 'a' corresponding to green-red axis in CIELAB. 

184 b : `NDArray` 

185 Chromaticity parameter 'b' corresponding to blue-yellow axis in CIELAB. 

186 saturation : `float`, optional 

187 Desired saturation level for the output colors. Defaults to 1. 

188 maxChroma : `float`, optional 

189 Maximum chroma value allowed for any color. Defaults to 50. 

190 

191 Returns 

192 ------- 

193 new_a : NDArray 

194 New a values representing the adjusted chromaticity. 

195 new_b : NDArray 

196 New b values representing the adjusted chromaticity. 

197 """ 

198 # Calculate the square of the chroma, which is the distance from origin in 

199 # the a-b plane. 

200 chroma1_2 = a**2 + b**2 

201 chroma1 = np.sqrt(chroma1_2) 

202 

203 # Calculate the hue angle, taking the absolute value to ensure non-negative 

204 # angle representation. 

205 chromaMask = chroma1 == 0 

206 chroma1[chromaMask] = 1 

207 sinHue = b / chroma1 

208 cosHue = a / chroma1 

209 sinHue[chromaMask] = 0 

210 cosHue[chromaMask] = 0 

211 

212 # Compute a divisor for saturation calculation, adding 1 to avoid division 

213 # by zero. 

214 div = chroma1_2 + oldLum**2 

215 div[div <= 0] = 1 

216 

217 # Calculate the square of the new chroma based on desired saturation 

218 sat_original_2 = chroma1_2 / div 

219 chroma2_2 = saturation * sat_original_2 * luminance**2 / (1 - sat_original_2) 

220 

221 # Cap the chroma to avoid excessive values that are visually unrealistic 

222 chroma2_2[chroma2_2 > maxChroma**2] = maxChroma**2 

223 

224 # Compute new 'a' values using the square root of adjusted chroma and 

225 # considering hue direction. 

226 chroma2 = np.sqrt(chroma2_2) 

227 new_a = chroma2 * cosHue 

228 

229 # Compute new 'b' values using the root of the adjusted chroma and hue 

230 # direction. 

231 new_b = chroma2 * sinHue 

232 

233 return new_a, new_b 

234 

235 

236def fixOutOfGamutColors( 

237 Lab: NDArray, colourspace: str = "Display P3", gamutMethod: Literal["mapping", "inpaint"] = "inpaint" 

238) -> NDArray: 

239 """Remap colors that fall outside an RGB color gamut back into it. 

240 

241 This function modifies the input Lab array in-place for memory reasons. 

242 

243 Parameters 

244 ---------- 

245 Lab : `NDArray` 

246 A NxMX3 array that contains data in the Lab colorspace. 

247 colourspace : `str`, optional 

248 The target colourspace to map outlying pixels into. This must 

249 correspond to an RGB colourspace understood by the colour-science 

250 python package. 

251 gamut_method : `str`, optional 

252 This determines what algorithm will be used to map out of gamut 

253 colors. Must be one of ``mapping`` or ``inpaint``. 

254 """ 

255 # Convert back into the CIE XYZ colourspace. 

256 xyz_prime = colour.Oklab_to_XYZ(Lab) 

257 

258 # And then back to the specified RGB colourspace. 

259 rgb_prime = colour.XYZ_to_RGB(xyz_prime, colourspace=colourspace) 

260 

261 # Determine if there are any out of bounds pixels 

262 outOfBounds = np.bitwise_or( 

263 np.bitwise_or(rgb_prime[:, :, 0] > 1, rgb_prime[:, :, 1] > 1), rgb_prime[:, :, 2] > 1 

264 ) 

265 

266 # If all pixels are in bounds, return immediately. 

267 if not np.any(outOfBounds): 

268 logging.info("There are no out of gamut pixels.") 

269 return rgb_prime 

270 

271 logging.info("There are out of gamut pixels, remapping colors") 

272 match gamutMethod: 

273 case "inpaint": 

274 results = inpaint_biharmonic(rgb_prime, outOfBounds, channel_axis=-1) 

275 case "mapping": 

276 results = fixGamutOK(Lab[outOfBounds]) 

277 Lab[outOfBounds] = results 

278 results = colour.XYZ_to_RGB(colour.Oklab_to_XYZ(Lab), colourspace=colourspace) 

279 case _: 

280 raise ValueError(f"gamut correction {gamutMethod} is not supported") 

281 

282 logging.debug(f"The total number of remapped pixels is: {np.sum(outOfBounds)}") 

283 return results 

284 

285 

286def _fuseExposure(images, sigma=0.2, maxLevel=3): 

287 weights = np.zeros((len(images), *images[0].shape[:2])) 

288 for i, image in enumerate(images): 

289 exposure = np.exp(-((image[:, :, 0] - 0.5) ** 2) / (2 * sigma)) 

290 

291 weights[i, :, :] = exposure 

292 norm = np.sum(weights, axis=0) 

293 np.divide(weights, norm, out=weights) 

294 

295 # loop over each image again to build pyramids 

296 g_pyr = [] 

297 l_pyr = [] 

298 maxImageLevel = int(np.min(np.log2(images[0].shape[:2]))) 

299 if maxLevel is None: 

300 maxLevel = maxImageLevel 

301 if maxImageLevel < maxLevel: 

302 raise ValueError( 

303 f"The supplied max level {maxLevel} is is greater than the max of the image: {maxImageLevel}" 

304 ) 

305 support = 1 << (maxLevel - 1) 

306 padY_amounts = levelPadder(image.shape[0] + support, maxLevel) 

307 padX_amounts = levelPadder(image.shape[1] + support, maxLevel) 

308 for image, weight in zip(images, weights): 

309 imagePadded = cv2.copyMakeBorder( 

310 image, *(0, support), *(0, support), cv2.BORDER_REPLICATE, None, None 

311 ).astype(image.dtype) 

312 weightPadded = cv2.copyMakeBorder( 

313 weight, *(0, support), *(0, support), cv2.BORDER_REPLICATE, None, None 

314 ).astype(image.dtype) 

315 

316 g_pyr.append(list(makeGaussianPyramid(weightPadded, padY_amounts, padX_amounts, None))) 

317 l_pyr.append(list(makeLapPyramid(imagePadded, padY_amounts, padX_amounts, None, None))) 

318 

319 # time to blend 

320 blended = [] 

321 for level in range(len(padY_amounts)): 

322 accumulate = np.zeros_like(l_pyr[0][level]) 

323 for img in range(len(g_pyr)): 

324 for i in range(3): 

325 accumulate[:, :, i] += l_pyr[img][level][:, :, i] * g_pyr[img][level] 

326 blended.append(accumulate) 

327 

328 # time to reconstruct 

329 output = blended[-1] 

330 for i in range(-2, -1 * len(blended) - 1, -1): 

331 upsampled = cv2.pyrUp(output) 

332 upsampled = upsampled[ 

333 : upsampled.shape[0] - 2 * padY_amounts[i + 1], : upsampled.shape[1] - 2 * padX_amounts[i + 1] 

334 ] 

335 output = blended[i] + upsampled 

336 return output[:-support, :-support] 

337 

338 

339def _handelLuminance( 

340 img: NDArray, 

341 scaleLum: Callable[..., NDArray] | None = latLum, 

342 scaleLumKWargs: Mapping | None = None, 

343 remapBounds: Callable | None = mapUpperBounds, 

344 remapBoundsKwargs: Mapping | None = None, 

345 doLocalContrast: bool = True, 

346 sigma: float = 30, 

347 highlights: float = -0.9, 

348 shadows: float = 0.5, 

349 clarity: float = 0.1, 

350 maxLevel: int | None = None, 

351 cieWhitePoint: tuple[float, float] = (0.28, 0.28), 

352 bracket: float = 1, 

353 psf: NDArray | None = None, 

354): 

355 # remap the bounds of the image if there is a function to do so. 

356 if remapBounds is not None: 

357 img = remapBounds(img, **(remapBoundsKwargs or {})) 

358 

359 # scale to the supplied bracket 

360 img /= bracket 

361 

362 # Convert the starting image into the OK L*a*b* color space. 

363 # https://en.wikipedia.org/wiki/Oklab_color_space 

364 

365 Lab = colour.XYZ_to_Oklab( 

366 colour.RGB_to_XYZ( 

367 img, 

368 colourspace="CIE RGB", 

369 illuminant=np.array(cieWhitePoint), 

370 chromatic_adaptation_transform="bradford", 

371 ) 

372 ) 

373 lum = Lab[:, :, 0] 

374 

375 # Enhance the contrast of the input image before mapping. 

376 if doLocalContrast: 

377 newLum = localContrast(lum, sigma, highlights, shadows, clarity=clarity, maxLevel=maxLevel) 

378 newLum = np.clip(newLum, 0, 1) 

379 else: 

380 newLum = lum 

381 

382 # Scale the luminance channel if possible. 

383 if scaleLum is not None: 

384 lRemapped = scaleLum(newLum, **(scaleLumKWargs or {})) 

385 else: 

386 lRemapped = newLum 

387 

388 if psf is not None: 

389 lRemapped = skimage.restoration.richardson_lucy(lRemapped, psf=psf, clip=False, num_iter=2) 

390 return lRemapped, Lab 

391 

392 

393def lsstRGB( 

394 rArray: NDArray, 

395 gArray: NDArray, 

396 bArray: NDArray, 

397 doLocalContrast: bool = True, 

398 scaleLum: Callable[..., NDArray] | None = latLum, 

399 scaleLumKWargs: Mapping | None = None, 

400 scaleColor: Callable[..., tuple[NDArray, NDArray]] | None = colorConstantSat, 

401 scaleColorKWargs: Mapping | None = None, 

402 remapBounds: Callable | None = mapUpperBounds, 

403 remapBoundsKwargs: Mapping | None = None, 

404 cieWhitePoint: tuple[float, float] = (0.28, 0.28), 

405 sigma: float = 30, 

406 highlights: float = -0.9, 

407 shadows: float = 0.5, 

408 clarity: float = 0.1, 

409 maxLevel: int | None = None, 

410 psf: NDArray | None = None, 

411 brackets: list[float] | None = None, 

412 doRemapGamut: bool = True, 

413 gamutMethod: Literal["mapping", "inpaint"] = "inpaint", 

414) -> NDArray: 

415 """Enhance the lightness and color preserving hue using perceptual methods. 

416 

417 Parameters 

418 ---------- 

419 rArray : `NDArray` 

420 The array used as the red channel 

421 gArray : `NDArray` 

422 The array used as the green channel 

423 bArray : `NDArray` 

424 The array used as the blue channel 

425 doLocalContrast: `bool` 

426 Apply local contrast enhancement algorithms to the luminance channel. 

427 scaleLum : `Callable` or `None` 

428 This is a callable that's passed the luminance values as well as 

429 any defined scaleLumKWargs, and should return a scaled luminance array 

430 the same shape as the input. Set to None for no scaling. 

431 scaleLumKWargs : `Mapping` or `None` 

432 Key word arguments that passed to the scaleLum function. 

433 scaleColor : `Callable` or `None` 

434 This is a callable that's passed the original luminance, the remapped 

435 luminance values, the a values for each pixel, and the b values for 

436 each pixel. The function is also passed any parameters defined in 

437 scaleColorKWargs. This function is responsible for scaling chroma 

438 values. This should return two arrays corresponding to the scaled a and 

439 b values. Set to None for no modification. 

440 scaleColorKWargs : `Mapping` or `None` 

441 Key word arguments passed to the scaleColor function. 

442 remapBounds : `Callable` or `None` 

443 This is a callable that should remaps the input arrays such that each of 

444 them fall within a zero to one range. This callable is given the 

445 initial image as well as any parameters defined in the remapBoundsKwargs 

446 parameter. Set to None for no remapping. 

447 remapBoundsKwargs : `Mapping` or None 

448 cieWhitePoint : `tuple` of `float`, `float` 

449 This is the white point of the input of the input arrays in CIE XY 

450 coordinates. Altering this affects the relative balance of colors 

451 in the input image, and therefore also the output image. 

452 sigma : `float` 

453 The scale over which local contrast considers edges real and not noise. 

454 highlights : `float` 

455 A parameter that controls how local contrast enhances or reduces 

456 highlights. Contrary to intuition, negative values increase highlights. 

457 shadows : `float` 

458 A parameter that controls how local contrast will deepen or reduce 

459 shadows. 

460 clarity : `float` 

461 A parameter that relates to the local contrast between highlights and 

462 shadow. 

463 maxLevel : `int` or `None` 

464 The maximum number of image pyramid levels to enhance the local contrast 

465 over. Each level has a spatial scale of roughly 2^(level) pixels. 

466 psf : `NDArray` or `None` 

467 If this parameter is an image of a PSF kernel the luminance channel is 

468 deconvolved with it. Set to None to skip deconvolution. 

469 brackets : `list` of `float` or `None` 

470 If a list brackets is supplied, an image will be generated at each of 

471 the brackets and the results will be used in exposure fusioning to 

472 increase the apparent dynamic range of the image. The image post bounds 

473 remapping will be divided by each of the values specified in this list, 

474 which can be used to create for instance an under, over, and ballanced 

475 expoisure. Theese will then be fusioned into a final single exposure 

476 selecting the proper elements from each of the images. 

477 doRemapGamut : `bool`, optional 

478 If this is `True` then any pixels which lay outside the representable 

479 color gamut after manipulation will be remapped to a "best" value 

480 which will be some compromise in hue, chroma, and lum. If this is 

481 `False` then the values will clip. This may be useful for 

482 seeing where problems in processing occur. 

483 

484 Returns 

485 ------- 

486 result : `NDArray` 

487 The brightness and color calibrated image. 

488 

489 Raises 

490 ------ 

491 ValueError 

492 Raised if the shapes of the input array don't match 

493 """ 

494 if rArray.shape != gArray.shape or rArray.shape != bArray.shape: 

495 raise ValueError("The shapes of all the input arrays must be the same") 

496 

497 # Construct a new image array in the proper byte ordering. 

498 img = np.empty((*rArray.shape, 3)) 

499 img[:, :, 0] = rArray 

500 img[:, :, 1] = gArray 

501 img[:, :, 2] = bArray 

502 # If there are nan's in the image there is no real option other than to 

503 # set them to zero or throw. 

504 img[np.isnan(img)] = 0 

505 

506 if not brackets: 

507 brackets = [1] 

508 

509 exposures = [] 

510 for bracket in brackets: 

511 tmp_lum, Lab = _handelLuminance( 

512 img, 

513 scaleLum, 

514 scaleLumKWargs=scaleLumKWargs, 

515 remapBounds=remapBounds, 

516 remapBoundsKwargs=remapBoundsKwargs, 

517 doLocalContrast=doLocalContrast, 

518 sigma=sigma, 

519 highlights=highlights, 

520 clarity=clarity, 

521 shadows=shadows, 

522 maxLevel=maxLevel, 

523 cieWhitePoint=cieWhitePoint, 

524 bracket=bracket, 

525 psf=psf, 

526 ) 

527 if scaleColor is not None: 

528 new_a, new_b = scaleColor( 

529 Lab[:, :, 0], tmp_lum, Lab[:, :, 1], Lab[:, :, 2], **(scaleColorKWargs or {}) 

530 ) 

531 # Replace the color information with the new scaled color information. 

532 Lab[:, :, 1] = new_a 

533 Lab[:, :, 2] = new_b 

534 # Replace the luminance information with the new scaled luminance information 

535 Lab[:, :, 0] = tmp_lum 

536 exposures.append(Lab) 

537 if len(brackets) > 1: 

538 Lab = _fuseExposure(exposures) 

539 

540 # Fix any colors that fall outside of the RGB colour gamut. 

541 if doRemapGamut: 

542 result = fixOutOfGamutColors(Lab, gamutMethod=gamutMethod) 

543 

544 # explicitly cut at 1 even though the mapping was to map colors 

545 # appropriately because the Z matrix transform can produce values greater 

546 # than 1 and is a known feature of the transform. 

547 result = np.clip(result, 0, 1) 

548 return result