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

77 statements  

« prev     ^ index     » next       coverage.py v7.13.5, created at 2026-04-23 08:45 +0000

1# This file is part of pipe_tasks. 

2# 

3# Developed for the LSST Data Management System. 

4# This product includes software developed by the LSST Project 

5# (https://www.lsst.org). 

6# See the COPYRIGHT file at the top-level directory of this distribution 

7# for details of code ownership. 

8# 

9# This program is free software: you can redistribute it and/or modify 

10# it under the terms of the GNU General Public License as published by 

11# the Free Software Foundation, either version 3 of the License, or 

12# (at your option) any later version. 

13# 

14# This program is distributed in the hope that it will be useful, 

15# but WITHOUT ANY WARRANTY; without even the implied warranty of 

16# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 

17# GNU General Public License for more details. 

18# 

19# You should have received a copy of the GNU General Public License 

20# along with this program. If not, see <https://www.gnu.org/licenses/>. 

21 

22__all__ = ("lsstRGB", "DEFAULT_FUNCTION") 

23 

24import logging 

25import numpy as np 

26import skimage 

27import time 

28 

29from enum import Enum, auto 

30 

31 

32from .types import ( 

33 FloatImagePlane, 

34 RGBImage, 

35 LocalContrastFunction, 

36 ScaleLumFunction, 

37 ScaleColorFunction, 

38 RemapBoundsFunction, 

39 BracketingFunction, 

40 GamutRemappingFunction, 

41) 

42 

43from ._functors import ( 

44 LocalContrastEnhancer, 

45 LumCompressor, 

46 ColorScaler, 

47 BoundsRemapper, 

48 ExposureBracketer, 

49 GamutFixer, 

50) 

51 

52 

53from lsst.rubinoxide import rgb 

54 

55logger = logging.getLogger(__name__) 

56 

57 

58class _SentinalDefault(Enum): 

59 DEFAULT_FUNCTION = auto() 

60 

61 

62DEFAULT_FUNCTION = _SentinalDefault.DEFAULT_FUNCTION 

63 

64 

65def lsstRGB( 

66 rArray: FloatImagePlane, 

67 gArray: FloatImagePlane, 

68 bArray: FloatImagePlane, 

69 local_contrast: LocalContrastFunction | None | _SentinalDefault = DEFAULT_FUNCTION, 

70 scale_lum: ScaleLumFunction | None | _SentinalDefault = DEFAULT_FUNCTION, 

71 scale_color: ScaleColorFunction | None | _SentinalDefault = DEFAULT_FUNCTION, 

72 remap_bounds: RemapBoundsFunction | None | _SentinalDefault = DEFAULT_FUNCTION, 

73 bracketing_function: BracketingFunction | None | _SentinalDefault = DEFAULT_FUNCTION, 

74 gamut_remapping_function: GamutRemappingFunction | None | _SentinalDefault = DEFAULT_FUNCTION, 

75 psf: FloatImagePlane | None = None, 

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

77) -> RGBImage: 

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

79 

80 Parameters 

81 ---------- 

82 rArray : `numpy.ndarray` 

83 The array used as the red channel 

84 gArray : `numpy.ndarray` 

85 The array used as the green channel 

86 bArray : `numpy.ndarray` 

87 The array used as the blue channel 

88 local_contrast : `LocalContrastFunction` or `None` 

89 This is a callable that's passed the luminance values, and is expected 

90 to do local contrast enhancement. Set to None to bypass. 

91 scale_lum : `ScaleLumFunction` or `None` 

92 This is a callable that's passed the luminance values and should 

93 return a scaled luminance array the same shape as the input. 

94 Set to None for no scaling. 

95 scale_color : `ScaleColorFunction` or `None` 

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

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

98 each pixel. This function is responsible for scaling chroma 

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

100 b values. Set to None for no modification. 

101 remap_bounds : `RemapBoundsFunction` or `None` 

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

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

104 initial image. Set to None for no remapping. If this is None 

105 input arrays should already be in 0-1 range. 

106 bracketing_function : `BracketingFunction` or `None` 

107 This is a callable that is passed the input luminance, and should 

108 create various exposure levels and then fuse them together. 

109 Set to None for no bracketing. 

110 gamut_remapping_function : `GamutRemappingFunction` or `None` 

111 This is a callable that is passed the final OkLab image. Its job 

112 is to detect and correct any pixel values that would fall outside 

113 an RGB P3 colorspace. Set to None for no fixes. 

114 psf : `FloatImagePlane` or `None` 

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

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

117 cieWhitePoint : `tuple` of `float`, `float` 

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

119 coordinates. Altering this affects the relative balance of colors 

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

121 

122 Returns 

123 ------- 

124 result : `RGBImage` 

125 The brightness and color calibrated image. 

126 

127 Raises 

128 ------ 

129 ValueError 

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

131 """ 

132 

133 # Default construct functors to be used as callables 

134 if local_contrast is DEFAULT_FUNCTION: 

135 local_contrast = LocalContrastEnhancer() 

136 if scale_lum is DEFAULT_FUNCTION: 

137 scale_lum = LumCompressor() 

138 if scale_color is DEFAULT_FUNCTION: 

139 scale_color = ColorScaler() 

140 if remap_bounds is DEFAULT_FUNCTION: 

141 remap_bounds = BoundsRemapper() 

142 if bracketing_function is DEFAULT_FUNCTION: 

143 bracketing_function = ExposureBracketer() 

144 if gamut_remapping_function is DEFAULT_FUNCTION: 

145 gamut_remapping_function = GamutFixer() 

146 

147 # Validate inputs 

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

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

150 

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

152 img: RGBImage = np.empty((*rArray.shape, 3)) 

153 img[:, :, 0] = rArray 

154 img[:, :, 1] = gArray 

155 img[:, :, 2] = bArray 

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

157 # set them to zero or throw. 

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

159 

160 logger.debug("Starting color processing pipeline") 

161 t1 = time.time() 

162 

163 if remap_bounds is not None: 

164 img = remap_bounds(img) 

165 logger.debug("doing remap took %.3fs", time.time() - t1) 

166 t1 = time.time() 

167 

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

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

170 Lab = rgb.RGB_to_Oklab(img, cieWhitePoint) 

171 logger.debug("lab conversion took %.3fs", time.time() - t1) 

172 t1 = time.time() 

173 lum = Lab[:, :, 0] 

174 

175 # potentially needed for remapping color, so save what it originally was 

176 lum_save = np.copy(lum) 

177 

178 if bracketing_function is not None: 

179 lum = bracketing_function(lum) 

180 logger.debug("bracketing took %.3fs", time.time() - t1) 

181 t1 = time.time() 

182 

183 if scale_lum is not None: 

184 lum = scale_lum(lum) 

185 logger.debug("lum scale took %.3fs", time.time() - t1) 

186 t1 = time.time() 

187 

188 if local_contrast is not None: 

189 lum = local_contrast(lum) 

190 logger.debug("local_contrast took %.3fs", time.time() - t1) 

191 t1 = time.time() 

192 

193 if psf is not None: 

194 lum = skimage.restoration.richardson_lucy(lum, psf=psf, clip=False, num_iter=2) 

195 logger.debug("psf took %.3fs", time.time() - t1) 

196 t1 = time.time() 

197 

198 if scale_color is not None: 

199 new_a, new_b = scale_color(lum_save, lum, Lab[..., 1], Lab[..., 2]) 

200 Lab[..., 1] = new_a 

201 Lab[..., 2] = new_b 

202 logger.debug("color correction took %.3fs", time.time() - t1) 

203 t1 = time.time() 

204 Lab[..., 0] = lum 

205 

206 # The target output profile whitepoint 

207 cie_white_point_d65 = (0.31272, 0.32903) 

208 if gamut_remapping_function is not None: 

209 result = gamut_remapping_function(Lab, cie_white_point_d65) 

210 logger.debug("gamut fixing took %.3fs", time.time() - t1) 

211 t1 = time.time() 

212 else: 

213 result = rgb.Oklab_to_RGB(np.ascontiguousarray(Lab), cie_white_point_d65) 

214 logger.debug("RGB conversion took %.3fs", time.time() - t1) 

215 t1 = time.time() 

216 

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

218 return result