Coverage for python / lsst / pipe / tasks / prettyPictureMaker / _localContrast.py: 8%

146 statements  

« prev     ^ index     » next       coverage.py v7.13.5, created at 2026-04-18 09:04 +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 

22from __future__ import annotations 

23 

24__all__ = ("localContrast",) 

25 

26import numpy as np 

27from numpy.typing import NDArray 

28import cv2 

29from numba import njit, prange 

30from numba.typed.typedlist import List 

31from collections.abc import Sequence 

32from itertools import cycle 

33 

34 

35@njit(fastmath=True, parallel=True, error_model="numpy", nogil=True) 

36def r( 

37 img: NDArray, out: NDArray, g: float, sigma: float, shadows: float, highlights: float, clarity: float 

38) -> NDArray: 

39 """ 

40 Apply a post-processing effect to an image using the specified parameters. 

41 

42 Parameters 

43 ---------- 

44 img : `numpy.ndarray` 

45 The input image array of shape (n_images, height, width). 

46 out : `numpy.ndarray` 

47 The output image array where the result will be stored. Should have the same shape as `img`. 

48 g : `float` 

49 A parameter for gamma correction. 

50 sigma : `float` 

51 Parameter that defines the scale at which a change should be considered an edge. 

52 shadows : `float` 

53 Shadow adjustment factor. 

54 highlights : `float` 

55 Highlight adjustment factor. Negative values INCREASE highlights. 

56 clarity : `float` 

57 Clarity adjustment factor. 

58 

59 Returns 

60 ------- 

61 result : `numpy.ndarray` 

62 The processed image array with the same shape as `out`. 

63 """ 

64 

65 h_s = (highlights, shadows) 

66 

67 # Iterate over each pixel in the image 

68 for i in prange(out.shape[0]): 

69 # Get the current image slice 

70 imgI = img[i] 

71 # Get the corresponding output slice 

72 outI = out[i] 

73 

74 # Iterate over each pixel in the image 

75 for j in prange(out.shape[1]): 

76 # Calculate the contrast adjusted by gamma correction 

77 c = imgI[j] - g 

78 # Determine the sign of the contrast adjustment 

79 s = np.sign(c) 

80 

81 # Compute the transformation term t based on the signed contrast 

82 t = s * c / (2.0 * sigma) 

83 # Clamp t to be within [0, 1] 

84 t = max(0, min(t, 1)) 

85 

86 t2 = t * t 

87 # Complement of t 

88 mt = 1.0 - t 

89 

90 # Determine the index based on the sign of c (either 0 or 1) 

91 index = np.uint8(np.bool_(1 + s)) 

92 

93 # Compute the final pixel value using the transformation and 

94 # additional terms for shadows/highlights and clarity 

95 val = g + s * sigma * 2 * mt * t + t2 * (s * sigma + s * sigma * h_s[index]) 

96 val = val + clarity * c * np.exp(-(c * c) / (2.0 * sigma * sigma / 3.0)) 

97 

98 # Assign the computed value to the output image 

99 outI[j] = val 

100 

101 return out 

102 

103 

104def makeGaussianPyramid( 

105 img: NDArray, padY: list[int], padX: list[int], out: List[NDArray] | None 

106) -> Sequence[NDArray]: 

107 """ 

108 Create a Gaussian Pyramid from an input image. 

109 

110 Parameters 

111 ---------- 

112 img : `numpy.ndarray` 

113 The input image, which will be processed to create the pyramid. 

114 padY : `list` of `int` 

115 List containing padding sizes along the Y-axis for each level of the pyramid. 

116 padX : `list` of `int` 

117 List containing padding sizes along the X-axis for each level of the pyramid. 

118 out : `~numba.typed.typedlist.List` of `NDArray` or `None` 

119 Optional list to store the output images of the pyramid levels. 

120 If None, a new list is created. 

121 

122 Returns 

123 ------- 

124 pyramid : `~collections.abc.sequence` of `numpy.ndarray` 

125 A sequence of images representing the Gaussian Pyramid. 

126 

127 Notes 

128 ----- 

129 - The function creates a padded version of the input image and then 

130 reduces its size using `cv2.pyrDown` to generate each level of the 

131 pyramid. 

132 - If 'out' is provided, it will be used to store the pyramid levels; 

133 otherwise, a new list is dynamically created. 

134 - Padding is applied only if specified by non-zero values in `padY` and 

135 `padX`. 

136 """ 

137 # Initialize the output pyramid list if not provided 

138 if out is None: 

139 pyramid = List() 

140 else: 

141 pyramid = out 

142 

143 # Apply padding only if needed, ensuring the type matches the input image 

144 if padY[0] or padX[0]: 

145 paddedImage = cv2.copyMakeBorder( 

146 img, *(0, padY[0]), *(0, padX[0]), cv2.BORDER_REPLICATE, None if out is None else pyramid[0], None 

147 ).astype(img.dtype) 

148 else: 

149 paddedImage = img 

150 

151 # Store the first level of the pyramid (padded image) 

152 if out is None: 

153 pyramid.append(paddedImage) 

154 else: 

155 # This might not be sound all the time, copy might be needed! 

156 # Update the first level in the provided list 

157 pyramid[0] = paddedImage 

158 

159 # Generate each subsequent level of the Gaussian Pyramid 

160 for i in range(1, len(padY)): 

161 if padY[i] or padX[i]: 

162 paddedImage = cv2.copyMakeBorder( 

163 paddedImage, *(0, padY[i]), *(0, padX[i]), cv2.BORDER_REPLICATE, None, None 

164 ).astype(img.dtype) 

165 # Downsample the image 

166 paddedImage = cv2.pyrDown(paddedImage, None if out is None else pyramid[i]) 

167 

168 # Append to the list if not provided externally 

169 if out is None: 

170 pyramid.append(paddedImage) 

171 return pyramid 

172 

173 

174def makeLapPyramid( 

175 img: NDArray, 

176 padY: list[int], 

177 padX: list[int], 

178 gaussOut: List[NDArray] | None, 

179 lapOut: List[NDArray] | None, 

180 upscratch: List[NDArray] | None = None, 

181) -> Sequence[NDArray]: 

182 """ 

183 Create a Laplacian pyramid from the input image. 

184 

185 This function constructs a Laplacian pyramid from the input image. It first 

186 generates a Gaussian pyramid and then, for each level (except the last), 

187 subtracts the upsampled version of the next lower level from the current 

188 level to obtain the Laplacian levels. If `lapOut` is None, it creates a 

189 new list to store the Laplacian pyramid; otherwise, it uses the provided 

190 `lapOut`. 

191 

192 Parameters 

193 ---------- 

194 img : `numpy.ndarray` 

195 The input image as a numpy array. 

196 padY : `list` of `int` 

197 List of padding sizes for rows (vertical padding). 

198 padX : `list` of `int` 

199 List of padding sizes for columns (horizontal padding). 

200 gaussOut : `~numba.typed.typedlist.List` of `numpy.ndarray` or None 

201 Preallocated storage for the output of the Gaussian pyramid function. 

202 If `None` new storage is allocated. 

203 lapOut : `~numba.typed.typedlist.List` of `numpy.ndarray` or None 

204 Preallocated for the output Laplacian pyramid. If None, a new 

205 `List` is created. 

206 upscratch : `~numba.typed.typedlist.List` of `numpy.ndarray`, optional 

207 List to store intermediate results of pyramids (default is None). 

208 

209 Returns 

210 ------- 

211 results : `~collections.abc.sequence` of `numpy.ndarray` 

212 The Laplacian pyramid as a sequence of numpy arrays. 

213 

214 """ 

215 pyramid = makeGaussianPyramid(img, padY, padX, gaussOut) 

216 if lapOut is None: 

217 lapPyramid = List() 

218 else: 

219 lapPyramid = lapOut 

220 for i in range(len(pyramid) - 1): 

221 upsampled = cv2.pyrUp(pyramid[i + 1], None if upscratch is None else upscratch[i + 1]) 

222 if padY[i + 1] or padX[i + 1]: 

223 upsampled = upsampled[ 

224 : upsampled.shape[0] - 2 * padY[i + 1], : upsampled.shape[1] - 2 * padX[i + 1] 

225 ] 

226 if lapOut is None: 

227 lapPyramid.append(pyramid[i] - upsampled) 

228 else: 

229 cv2.subtract(pyramid[i], upsampled, dst=lapPyramid[i]) 

230 if lapOut is None: 

231 lapPyramid.append(pyramid[-1]) 

232 else: 

233 lapPyramid[-1][:, :] = pyramid[-1] 

234 return lapPyramid 

235 

236 

237@njit(fastmath=True, parallel=True, error_model="numpy", nogil=True) 

238def _calculateOutput( 

239 out: List[NDArray], 

240 pyramid: List[NDArray], 

241 gamma: NDArray, 

242 pyramidVectorsBottom: List[NDArray], 

243 pyramidVectorsTop: List[NDArray], 

244): 

245 """ 

246 Computes the output by interpolating between basis vectors at each pixel in 

247 a Gaussian pyramid. 

248 

249 The function iterates over each pixel in the Gaussian pyramids 

250 and interpolates between the corresponding basis vectors from 

251 `pyramidVectorsBottom` and `pyramidVectorsTop`. If a pixel value is outside 

252 the range defined by gamma, it skips interpolation. 

253 

254 Parameters 

255 ---------- 

256 out : `~numba.typed.typedlist.List` of `numpy.ndarray` 

257 A list of numpy arrays representing the output image pyramids. 

258 pyramid : `~numba.typed.typedlist.List` of `numpy.ndarray` 

259 A list of numpy arrays representing the Gaussian pyramids. 

260 gamma : `numpy.ndarray` 

261 A numpy array containing the range for pixel values to be considered in 

262 the interpolation. 

263 pyramidVectorsBottom : `~numba.typed.typedlist.List` of `numpy.ndarray` 

264 A list of numpy arrays representing the basis vectors at the bottom 

265 level of each pyramid layer. 

266 pyramidVectorsTop : `~numba.typed.typedlist.List` of `numpy.ndarray` 

267 A list of numpy arrays representing the basis vectors at the top level 

268 of each pyramid layer. 

269 

270 """ 

271 # loop over each pixel in the gaussian pyramid 

272 for level in prange(0, len(pyramid)): 

273 yshape = pyramid[level].shape[0] 

274 xshape = pyramid[level].shape[1] 

275 plevel = pyramid[level] 

276 outlevel = out[level] 

277 basisBottom = pyramidVectorsBottom[level] 

278 basisTop = pyramidVectorsTop[level] 

279 for y in prange(yshape): 

280 plevelY = plevel[y] 

281 outLevelY = outlevel[y] 

282 basisBottomY = basisBottom[y] 

283 basisTopY = basisTop[y] 

284 for x in prange(xshape): 

285 val = plevelY[x] 

286 if not (val >= gamma[0] and val < gamma[1]): 

287 continue 

288 a = (plevelY[x] - gamma[0]) / (gamma[1] - gamma[0]) 

289 outLevelY[x] = (1 - a) * basisBottomY[x] + a * basisTopY[x] 

290 

291 

292def levelPadder(numb: int, levels: int) -> list[int]: 

293 """Determine if each level of transform will need to be padded by one to 

294 ensure the level is divisible by two. 

295 

296 Parameters 

297 ---------- 

298 numb : `int` 

299 The size of the input dimension 

300 levels : `int` 

301 The number of times the dimensions will be reduced by a factor of two 

302 

303 Returns 

304 ------- 

305 pads : `list` of `int` 

306 A list where the entries are either zero or one depending on if the 

307 size will need pad to be a power of two. 

308 

309 """ 

310 pads = [] 

311 if numb % 2 != 0: 

312 pads.append(1) 

313 numb += 1 

314 else: 

315 pads.append(0) 

316 for _ in range(levels): 

317 numb /= 2 

318 if numb % 2 != 0: 

319 pads.append(1) 

320 numb += 1 

321 else: 

322 pads.append(0) 

323 return pads 

324 

325 

326def localContrast( 

327 image: NDArray, 

328 sigma: float, 

329 highlights: float = -0.9, 

330 shadows: float = 0.4, 

331 clarity: float = 0.15, 

332 maxLevel: int | None = None, 

333 numGamma: int = 20, 

334 skipLevels: int = 0, 

335) -> NDArray: 

336 """Enhance the local contrast of an input image. 

337 

338 Parameters 

339 ---------- 

340 image : `numpy.ndarray` 

341 Two dimensional numpy array representing the image to have contrast 

342 increased. Data must be in the range 0 to 1. 

343 sigma : `float` 

344 The scale over which edges are considered real and not noise. 

345 highlights : `float` 

346 A parameter that controls how highlights are enhanced or reduced, 

347 contrary to intuition, negative values increase highlights. 

348 shadows : `float` 

349 A parameter that controls how shadows are deepened. 

350 clarity : `float` 

351 A parameter that relates to the contrast between highlights and 

352 shadow. 

353 maxLevel : `int` or `None` 

354 The maximum number of image pyramid levels to enhance the contrast over. 

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

356 numGamma : `int` 

357 This is an optimization parameter. This algorithm divides up contrast 

358 space into a certain number of bins over which the expensive computation 

359 is done. Contrast values in the image which fall between two of these 

360 values are interpolated to get the outcome. The higher the numGamma, 

361 the smoother the image is post contrast enhancement, though above 

362 some number there is no discernible difference. 

363 skipLevels : `int` 

364 When calculating the local contrast skip the specified number of levels 

365 starting at the lowest level. 

366 

367 Returns 

368 ------- 

369 image : `numpy.ndarray` 

370 Two dimensional numpy array of the input image with increased local 

371 contrast. 

372 

373 Raises 

374 ------ 

375 ValueError 

376 Raised if the max level to enhance to is greater than the image 

377 supports. 

378 

379 Notes 

380 ----- 

381 This function, and its supporting functions, spiritually implement the 

382 algorithm outlined at 

383 https://people.csail.mit.edu/sparis/publi/2011/siggraph/ 

384 titled "Local Laplacian Filters: Edge-aware Image Processing with Laplacian 

385 Pyramid". This is not a 1:1 implementation, it's optimized for the 

386 python language and runtime performance. Most notably it transforms only 

387 certain levels and linearly interpolates to find other values. This 

388 implementation is inspired by the one done in the darktable image editor: 

389 https://www.darktable.org/2017/11/local-laplacian-pyramids/. None of the 

390 code is in common, nor is the implementation 1:1, but reading the original 

391 paper and the darktable implementation gives more info about this function. 

392 Specifically some variable names follow the paper/other implementation, 

393 and may be confusing when viewed without that context. 

394 

395 """ 

396 # ensure the supplied values are floats, and not ints 

397 highlights = float(highlights) 

398 shadows = float(shadows) 

399 clarity = float(clarity) 

400 

401 # Determine the maximum level over which the image will be enhanced 

402 # and the amount of padding that will be needed to be added to the 

403 # image. 

404 maxImageLevel = int(np.min(np.log2(image.shape))) 

405 if maxLevel is None: 

406 maxLevel = maxImageLevel 

407 if maxImageLevel < maxLevel: 

408 raise ValueError( 

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

410 ) 

411 support = 1 << (maxLevel - 1) 

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

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

414 imagePadded = cv2.copyMakeBorder( 

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

416 ).astype(image.dtype) 

417 

418 # build a list of intensities 

419 gamma = np.linspace(0, 1, numGamma) 

420 

421 # make gaussian pyramid 

422 pyramid = makeGaussianPyramid(imagePadded, padY_amounts, padX_amounts, None) 

423 base_lap_pyr = makeLapPyramid(imagePadded, padY_amounts, padX_amounts, None, None, None) 

424 

425 finalPyramid = List() 

426 for sample in pyramid[:-1]: 

427 finalPyramid.append(np.zeros_like(sample)) 

428 finalPyramid.append(pyramid[-1]) 

429 

430 # make a working array for gaussian pyramid in Lap 

431 # make two working arrays for laplace as the true value is interpolated 

432 # between the endpoints. 

433 # This prevents needing re-allocations which can be time consuming. 

434 tmpGauss = List() 

435 tmpLap1 = List() 

436 tmpLap2 = List() 

437 upscratch = List() 

438 for i, sample in enumerate(pyramid): 

439 tmpGauss.append(np.empty_like(sample)) 

440 tmpLap1.append(np.empty_like(sample)) 

441 tmpLap2.append(np.empty_like(sample)) 

442 if i == 0: 

443 upscratch.append(np.empty((0, 0), dtype=image.dtype)) 

444 continue 

445 upscratch.append(np.empty((sample.shape[0] * 2, sample.shape[1] * 2), dtype=image.dtype)) 

446 # cycle between the endpoints, because there is no reason to recalculate both 

447 # endpoints as only one changes for each bin. 

448 cycler = iter(cycle((tmpLap1, tmpLap2))) 

449 # allocate temporary arrays to use for each bin 

450 outCycle = iter(cycle((np.copy(imagePadded), np.copy(imagePadded)))) 

451 prevImg = r( 

452 imagePadded, next(outCycle), gamma[0], sigma, shadows=shadows, highlights=highlights, clarity=clarity 

453 ) 

454 prevLapPyr = makeLapPyramid( 

455 prevImg, padY_amounts, padX_amounts, tmpGauss, next(cycler), upscratch=upscratch 

456 ) 

457 

458 for value in range(1, len(gamma)): 

459 pyramidVectors = List() 

460 pyramidVectors.append(prevLapPyr) 

461 newImg = r( 

462 imagePadded, 

463 next(outCycle), 

464 gamma[value], 

465 sigma, 

466 shadows=shadows, 

467 highlights=highlights, 

468 clarity=clarity, 

469 ) 

470 prevLapPyr = makeLapPyramid( 

471 newImg, padY_amounts, padX_amounts, tmpGauss, next(cycler), upscratch=upscratch 

472 ) 

473 pyramidVectors.append(prevLapPyr) 

474 

475 _calculateOutput( 

476 finalPyramid, 

477 pyramid, 

478 np.array((gamma[value - 1], gamma[value])), 

479 pyramidVectors[0], 

480 pyramidVectors[1], 

481 ) 

482 del pyramidVectors 

483 if skipLevels: 

484 for i in range(skipLevels): 

485 finalPyramid[i] = base_lap_pyr[i] 

486 

487 # time to reconstruct 

488 output = finalPyramid[-1] 

489 for i in range(-2, -1 * len(finalPyramid) - 1, -1): 

490 upsampled = cv2.pyrUp(output) 

491 upsampled = upsampled[ 

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

493 ] 

494 output = finalPyramid[i] + upsampled 

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