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")
3import logging
4import numpy as np
5import skimage
6import colour
7import cv2
8from skimage.restoration import inpaint_biharmonic
10from ._localContrast import localContrast, makeGaussianPyramid, makeLapPyramid, levelPadder
11from lsst.cpputils import fixGamutOK
13from numpy.typing import NDArray
14from typing import Callable, Mapping, Literal
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.
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.
63 Returns:
64 luminance : `NDArray`
65 The stretched luminosity data.
66 """
68 # De-noise the input image using wavelet de-noising.
69 if doDenoise:
70 values = skimage.restoration.denoise_wavelet(values)
71 values = abs(values)
73 # Scale values from 0-1 to 0-100 as various algorithm expect that range.
74 values *= 100
76 # Calculate the slope for arcsinh transformation based on Q and stretch
77 # parameters.
78 slope = 0.1 * 100 / np.arcsinh(0.1 * Q)
80 # Apply the modified luminosity transformation using arcsinh function.
81 soften = Q / stretch
82 intensities = np.arcsinh((abs(values) * soften + floor) * slope)
84 # Always normalize by what the original max value (100) scales to.
85 maximum_intensity = np.arcsinh((100 * soften + floor) * slope)
87 intensities /= maximum_intensity
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)
93 np.clip(intensities, 0, max, out=intensities)
95 # Reset the input array.
96 values /= 100
98 return intensities
101def mapUpperBounds(img: NDArray, quant: float = 0.9, absMax: float | None = None) -> NDArray:
102 """Bound images to a range between zero and one.
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.
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.
127 Returns
128 -------
129 image : `NDArray`
130 The result of the remapping process
131 """
132 if np.max(img) == 1:
133 return img
135 r = img[:, :, 0]
136 g = img[:, :, 1]
137 b = img[:, :, 2]
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))
144 if absMax is not None:
145 scale = absMax
146 else:
147 scale = turnover * quant
149 image = np.empty(img.shape)
150 image[:, :, 0] = r / scale
151 image[:, :, 1] = g / scale
152 image[:, :, 2] = b / scale
154 # Clip values that exceed the bound to ensure all values are within [0, absMax]
155 np.clip(image, 0, 1, out=image)
157 return image
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.
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.
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.
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)
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
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
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)
221 # Cap the chroma to avoid excessive values that are visually unrealistic
222 chroma2_2[chroma2_2 > maxChroma**2] = maxChroma**2
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
229 # Compute new 'b' values using the root of the adjusted chroma and hue
230 # direction.
231 new_b = chroma2 * sinHue
233 return new_a, new_b
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.
241 This function modifies the input Lab array in-place for memory reasons.
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)
258 # And then back to the specified RGB colourspace.
259 rgb_prime = colour.XYZ_to_RGB(xyz_prime, colourspace=colourspace)
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 )
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
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")
282 logging.debug(f"The total number of remapped pixels is: {np.sum(outOfBounds)}")
283 return results
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))
291 weights[i, :, :] = exposure
292 norm = np.sum(weights, axis=0)
293 np.divide(weights, norm, out=weights)
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)
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)))
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)
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]
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 {}))
359 # scale to the supplied bracket
360 img /= bracket
362 # Convert the starting image into the OK L*a*b* color space.
363 # https://en.wikipedia.org/wiki/Oklab_color_space
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]
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
382 # Scale the luminance channel if possible.
383 if scaleLum is not None:
384 lRemapped = scaleLum(newLum, **(scaleLumKWargs or {}))
385 else:
386 lRemapped = newLum
388 if psf is not None:
389 lRemapped = skimage.restoration.richardson_lucy(lRemapped, psf=psf, clip=False, num_iter=2)
390 return lRemapped, Lab
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.
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.
484 Returns
485 -------
486 result : `NDArray`
487 The brightness and color calibrated image.
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")
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
506 if not brackets:
507 brackets = [1]
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)
540 # Fix any colors that fall outside of the RGB colour gamut.
541 if doRemapGamut:
542 result = fixOutOfGamutColors(Lab, gamutMethod=gamutMethod)
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