Coverage for python / lsst / pipe / tasks / prettyPictureMaker / _task.py: 23%

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__ = ( 

25 "ChannelRGBConfig", 

26 "PrettyPictureTask", 

27 "PrettyPictureConnections", 

28 "PrettyPictureConfig", 

29 "PrettyMosaicTask", 

30 "PrettyMosaicConnections", 

31 "PrettyMosaicConfig", 

32 "PrettyPictureBackgroundFixerConfig", 

33 "PrettyPictureBackgroundFixerTask", 

34 "PrettyPictureStarFixerConfig", 

35 "PrettyPictureStarFixerTask", 

36) 

37 

38from collections.abc import Iterable, Mapping 

39from lsst.afw.image import ExposureF 

40import numpy as np 

41from typing import TYPE_CHECKING, cast, Any 

42from lsst.skymap import BaseSkyMap 

43 

44from scipy.stats import norm, halfnorm, mode 

45from scipy.ndimage import binary_dilation 

46from scipy.interpolate import RBFInterpolator 

47from skimage.restoration import inpaint_biharmonic 

48 

49from lsst.daf.butler import Butler, DeferredDatasetHandle 

50from lsst.daf.butler import DatasetRef 

51from lsst.pex.config import Field, Config, ConfigDictField, ConfigField, ListField, ChoiceField 

52from lsst.pipe.base import ( 

53 PipelineTask, 

54 PipelineTaskConfig, 

55 PipelineTaskConnections, 

56 Struct, 

57 InMemoryDatasetHandle, 

58) 

59import cv2 

60 

61from lsst.pipe.base.connectionTypes import Input, Output 

62from lsst.geom import Box2I, Point2I, Extent2I 

63from lsst.afw.image import Exposure, Mask 

64 

65from ._plugins import plugins 

66from ._colorMapper import lsstRGB 

67 

68import tempfile 

69 

70 

71if TYPE_CHECKING: 

72 from numpy.typing import NDArray 

73 from lsst.pipe.base import QuantumContext, InputQuantizedConnection, OutputQuantizedConnection 

74 from lsst.skymap import TractInfo, PatchInfo 

75 

76 

77class PrettyPictureConnections( 

78 PipelineTaskConnections, 

79 dimensions={"tract", "patch", "skymap"}, 

80 defaultTemplates={"coaddTypeName": "deep"}, 

81): 

82 inputCoadds = Input( 

83 doc=( 

84 "Model of the static sky, used to find temporal artifacts. Typically a PSF-Matched, " 

85 "sigma-clipped coadd. Written if and only if assembleStaticSkyModel.doWrite=True" 

86 ), 

87 name="pretty_coadd", 

88 storageClass="ExposureF", 

89 dimensions=("tract", "patch", "skymap", "band"), 

90 multiple=True, 

91 ) 

92 

93 outputRGB = Output( 

94 doc="A RGB image created from the input data stored as a 3d array", 

95 name="rgb_picture_array", 

96 storageClass="NumpyArray", 

97 dimensions=("tract", "patch", "skymap"), 

98 ) 

99 

100 outputRGBMask = Output( 

101 doc="A Mask corresponding to the fused masks of the input channels", 

102 name="rgb_picture_mask", 

103 storageClass="Mask", 

104 dimensions=("tract", "patch", "skymap"), 

105 ) 

106 

107 

108class ChannelRGBConfig(Config): 

109 """This describes the rgb values of a given input channel. 

110 

111 For instance if this channel is red the values would be self.r = 1, 

112 self.g = 0, self.b = 0. If the channel was cyan the values would be 

113 self.r = 0, self.g = 1, self.b = 1. 

114 """ 

115 

116 r = Field[float](doc="The amount of red contained in this channel") 

117 g = Field[float](doc="The amount of green contained in this channel") 

118 b = Field[float](doc="The amount of blue contained in this channel") 

119 

120 

121class LumConfig(Config): 

122 """Configurations to control how luminance is mapped in the rgb code""" 

123 

124 stretch = Field[float](doc="The stretch of the luminance in asinh", default=400) 

125 max = Field[float](doc="The maximum allowed luminance on a 0 to 100 scale", default=85) 

126 floor = Field[float](doc="A scaling factor to apply to the luminance before asinh scaling", default=0.0) 

127 Q = Field[float](doc="softening parameter", default=0.7) 

128 highlight = Field[float]( 

129 doc="The value of highlights in scaling factor applied to post asinh streaching", default=1.0 

130 ) 

131 shadow = Field[float]( 

132 doc="The value of shadows in scaling factor applied to post asinh streaching", default=0.0 

133 ) 

134 midtone = Field[float]( 

135 doc="The value of midtone in scaling factor applied to post asinh streaching", default=0.5 

136 ) 

137 

138 

139class LocalContrastConfig(Config): 

140 """Configuration to control local contrast enhancement of the luminance 

141 channel.""" 

142 

143 doLocalContrast = Field[bool]( 

144 doc="Apply local contrast enhancements to the luminance channel", default=True 

145 ) 

146 highlights = Field[float](doc="Adjustment factor for the highlights", default=-0.9) 

147 shadows = Field[float](doc="Adjustment factor for the shadows", default=0.5) 

148 clarity = Field[float](doc="Amount of clarity to apply to contrast modification", default=0.1) 

149 sigma = Field[float]( 

150 doc="The scale size of what is considered local in the contrast enhancement", default=30 

151 ) 

152 maxLevel = Field[int]( 

153 doc="The maximum number of scales the contrast should be enhanced over, if None then all", 

154 default=4, 

155 optional=True, 

156 ) 

157 

158 

159class ScaleColorConfig(Config): 

160 """Controls color scaling in the RGB generation process.""" 

161 

162 saturation = Field[float]( 

163 doc=( 

164 "The overall saturation factor with the scaled luminance between zero and one. " 

165 "A value of one is not recommended as it makes bright pixels very saturated" 

166 ), 

167 default=0.5, 

168 ) 

169 maxChroma = Field[float]( 

170 doc=( 

171 "The maximum chromaticity in the CIELCh color space, large " 

172 "values will cause bright pixels to fall outside the RGB gamut." 

173 ), 

174 default=50.0, 

175 ) 

176 

177 

178class RemapBoundsConfig(Config): 

179 """Remaps input images to a known range of values. 

180 

181 Often input images are not mapped to any defined range of values 

182 (for instance if they are in count units). This controls how the units of 

183 and image are mapped to a zero to one range by determining an upper 

184 bound. 

185 """ 

186 

187 quant = Field[float]( 

188 doc=( 

189 "The maximum values of each of the three channels will be multiplied by this factor to " 

190 "determine the maximum flux of the image, values larger than this quantity will be clipped." 

191 ), 

192 default=0.8, 

193 ) 

194 absMax = Field[float]( 

195 doc="Instead of determining the maximum value from the image, use this fixed value instead", 

196 default=220, 

197 optional=True, 

198 ) 

199 

200 

201class PrettyPictureConfig(PipelineTaskConfig, pipelineConnections=PrettyPictureConnections): 

202 channelConfig = ConfigDictField( 

203 doc="A dictionary that maps band names to their rgb channel configurations", 

204 keytype=str, 

205 itemtype=ChannelRGBConfig, 

206 default={}, 

207 ) 

208 imageRemappingConfig = ConfigField[RemapBoundsConfig]( 

209 doc="Configuration controlling channel normalization process" 

210 ) 

211 luminanceConfig = ConfigField[LumConfig]( 

212 doc="Configuration for the luminance scaling when making an RGB image" 

213 ) 

214 localContrastConfig = ConfigField[LocalContrastConfig]( 

215 doc="Configuration controlling the local contrast correction in RGB image production" 

216 ) 

217 colorConfig = ConfigField[ScaleColorConfig]( 

218 doc="Configuration to control the color scaling process in RGB image production" 

219 ) 

220 cieWhitePoint = ListField[float]( 

221 doc="The white point of the input arrays in ciexz coordinates", maxLength=2, default=[0.28, 0.28] 

222 ) 

223 arrayType = ChoiceField[str]( 

224 doc="The dataset type for the output image array", 

225 default="uint8", 

226 allowed={ 

227 "uint8": "Use 8 bit arrays, 255 max", 

228 "uint16": "Use 16 bit arrays, 65535 max", 

229 "half": "Use 16 bit float arrays, 1 max", 

230 "float": "Use 32 bit float arrays, 1 max", 

231 }, 

232 ) 

233 doPSFDeconcovlve = Field[bool]( 

234 doc="Use the PSF in a richardson lucy deconvolution on the luminance channel.", default=True 

235 ) 

236 exposureBrackets = ListField[float]( 

237 doc=( 

238 "Exposure scaling factors used in creating multiple exposures with different scalings which will " 

239 "then be fused into a final image" 

240 ), 

241 optional=True, 

242 default=[1.25, 1, 0.75], 

243 ) 

244 doRemapGamut = Field[bool]( 

245 doc="Apply a color correction to unrepresentable colors, if false they will clip", default=True 

246 ) 

247 gamutMethod = ChoiceField[str]( 

248 doc="If doRemapGamut is True this determines the method", 

249 default="inpaint", 

250 allowed={ 

251 "mapping": "Use a mapping function", 

252 "inpaint": "Use surrounding pixels to determine likely value", 

253 }, 

254 ) 

255 recenterNoise = Field[float]( 

256 doc="Recenter the noise away from zero. Supplied value is in units of sigma", 

257 optional=True, 

258 default=None, 

259 ) 

260 noiseSearchThreshold = Field[float]( 

261 doc=( 

262 "Flux threshold below which most flux will be considered noise, used to estimate noise properties" 

263 ), 

264 default=10, 

265 ) 

266 

267 def setDefaults(self): 

268 self.channelConfig["i"] = ChannelRGBConfig(r=1, g=0, b=0) 

269 self.channelConfig["r"] = ChannelRGBConfig(r=0, g=1, b=0) 

270 self.channelConfig["g"] = ChannelRGBConfig(r=0, g=0, b=1) 

271 return super().setDefaults() 

272 

273 

274class PrettyPictureTask(PipelineTask): 

275 """Turns inputs into an RGB image.""" 

276 

277 _DefaultName = "prettyPictureTask" 

278 ConfigClass = PrettyPictureConfig 

279 

280 config: ConfigClass 

281 

282 def _find_normal_stats(self, array): 

283 """Calculate standard deviation from negative values using half-normal distribution. 

284 

285 Parameters 

286 ---------- 

287 array : `numpy.array` 

288 Input array of numerical values. 

289 

290 Returns 

291 ------- 

292 mean : `float` 

293 The central moment of the distribution 

294 sigma : `float` 

295 Estimated standard deviation from negative values. Returns np.inf if: 

296 - No negative values exist in the array 

297 - Half-normal fitting fails 

298 """ 

299 # Extract negative values efficiently 

300 values_noise = array[array < self.config.noiseSearchThreshold] 

301 

302 # find the mode 

303 center = mode(np.round(values_noise, 2)).mode 

304 

305 # extract the negative values 

306 values_neg = array[array < center] 

307 

308 # Return infinity if no negative values found 

309 if values_neg.size == 0: 

310 return 0, np.inf 

311 

312 try: 

313 # Fit half-normal distribution to absolute negative values 

314 mu, sigma = halfnorm.fit(np.abs(values_neg)) 

315 except (ValueError, RuntimeError): 

316 # Handle fitting failures (e.g., constant data, optimization issues) 

317 return 0, np.inf 

318 

319 return center, sigma 

320 

321 def _match_sigmas_and_recenter(self, *arrays, factor=1): 

322 """Scale array values to match minimum standard deviation across arrays 

323 and recenter noise. 

324 

325 Adjusts values below each array's sigma by scaling and shifting them to 

326 align with the minimum sigma value across all input arrays. This operates 

327 in-place for efficiency. 

328 

329 Parameters 

330 ---------- 

331 *arrays : any number of `numpy.array` 

332 Variable number of input arrays to process. 

333 factor : float, optional 

334 Scaling factor for adjustments (default: 1). 

335 

336 """ 

337 # Calculate standard deviations for all arrays 

338 sigmas = [] 

339 mus = [] 

340 for arr in arrays: 

341 m, s = self._find_normal_stats(arr) 

342 mus.append(m) 

343 sigmas.append(s) 

344 mus = np.array(mus) 

345 sigmas = np.array(sigmas) 

346 

347 # If no sigmas could be determined, return the original 

348 # arrays. 

349 if not np.any(np.isfinite(sigmas)): 

350 return 

351 

352 min_sig = np.min(sigmas) 

353 

354 for mu, sigma, array in zip(mus, sigmas, arrays): 

355 # Identify values below the array's sigma threshold 

356 lower_pos = (array - mu) < sigma 

357 

358 # Skip processing if sigma is invalid 

359 if not np.isfinite(sigma): 

360 continue 

361 

362 # Calculate scaling ratio relative to minimum sigma 

363 sigma_ratio = min_sig / sigma 

364 

365 # Apply adjustment to qualifying values 

366 array[lower_pos] = (array[lower_pos] - mu) * sigma_ratio + min_sig * factor 

367 

368 def run(self, images: Mapping[str, Exposure]) -> Struct: 

369 """Turns the input arguments in arguments into an RGB array. 

370 

371 Parameters 

372 ---------- 

373 images : `Mapping` of `str` to `Exposure` 

374 A mapping of input images and the band they correspond to. 

375 

376 Returns 

377 ------- 

378 result : `Struct` 

379 A struct with the corresponding RGB image, and mask used in 

380 RGB image construction. The struct will have the attributes 

381 outputRGBImage and outputRGBMask. Each of the outputs will 

382 be a `NDarray` object. 

383 

384 Notes 

385 ----- 

386 Construction of input images are made easier by use of the 

387 makeInputsFrom* methods. 

388 """ 

389 channels = {} 

390 shape = (0, 0) 

391 jointMask: None | NDArray = None 

392 maskDict: Mapping[str, int] = {} 

393 doJointMaskInit = False 

394 if jointMask is None: 

395 doJointMask = True 

396 doJointMaskInit = True 

397 for channel, imageExposure in images.items(): 

398 imageArray = imageExposure.image.array 

399 # run all the plugins designed for array based interaction 

400 for plug in plugins.channel(): 

401 imageArray = plug( 

402 imageArray, imageExposure.mask.array, imageExposure.mask.getMaskPlaneDict(), self.config 

403 ).astype(np.float32) 

404 channels[channel] = imageArray 

405 # These operations are trivial look-ups and don't matter if they 

406 # happen in each loop. 

407 shape = imageArray.shape 

408 maskDict = imageExposure.mask.getMaskPlaneDict() 

409 if doJointMaskInit: 

410 jointMask = np.zeros(shape, dtype=imageExposure.mask.dtype) 

411 doJointMaskInit = False 

412 if doJointMask: 

413 jointMask |= imageExposure.mask.array 

414 

415 # mix the images to RGB 

416 imageRArray = np.zeros(shape, dtype=np.float32) 

417 imageGArray = np.zeros(shape, dtype=np.float32) 

418 imageBArray = np.zeros(shape, dtype=np.float32) 

419 

420 for band, image in channels.items(): 

421 mix = self.config.channelConfig[band] 

422 if mix.r: 

423 imageRArray += mix.r * image 

424 if mix.g: 

425 imageGArray += mix.g * image 

426 if mix.b: 

427 imageBArray += mix.b * image 

428 

429 exposure = next(iter(images.values())) 

430 box: Box2I = exposure.getBBox() 

431 boxCenter = box.getCenter() 

432 try: 

433 psf = exposure.psf.computeImage(boxCenter).array 

434 except Exception: 

435 psf = None 

436 

437 if self.config.recenterNoise: 

438 self._match_sigmas_and_recenter( 

439 imageRArray, imageGArray, imageBArray, factor=self.config.recenterNoise 

440 ) 

441 

442 # assert for typing reasons 

443 assert jointMask is not None 

444 # Run any image level correction plugins 

445 colorImage = np.zeros((*imageRArray.shape, 3)) 

446 colorImage[:, :, 0] = imageRArray 

447 colorImage[:, :, 1] = imageGArray 

448 colorImage[:, :, 2] = imageBArray 

449 for plug in plugins.partial(): 

450 colorImage = plug(colorImage, jointMask, maskDict, self.config) 

451 

452 # Ignore type because Exposures do in fact have a bbox, but it's c++ 

453 # and not typed. 

454 colorImage = lsstRGB( 

455 colorImage[:, :, 0], 

456 colorImage[:, :, 1], 

457 colorImage[:, :, 2], 

458 scaleLumKWargs=self.config.luminanceConfig.toDict(), 

459 remapBoundsKwargs=self.config.imageRemappingConfig.toDict(), 

460 scaleColorKWargs=self.config.colorConfig.toDict(), 

461 **(self.config.localContrastConfig.toDict()), 

462 cieWhitePoint=tuple(self.config.cieWhitePoint), # type: ignore 

463 psf=psf if self.config.doPSFDeconcovlve else None, 

464 brackets=list(self.config.exposureBrackets) if self.config.exposureBrackets else None, 

465 doRemapGamut=self.config.doRemapGamut, 

466 gamutMethod=self.config.gamutMethod, 

467 ) 

468 

469 # Find the dataset type and thus the maximum values as well 

470 maxVal: int | float 

471 match self.config.arrayType: 

472 case "uint8": 

473 dtype = np.uint8 

474 maxVal = 255 

475 case "uint16": 

476 dtype = np.uint16 

477 maxVal = 65535 

478 case "half": 

479 dtype = np.half 

480 maxVal = 1.0 

481 case "float": 

482 dtype = np.float32 

483 maxVal = 1.0 

484 case _: 

485 assert True, "This code path should be unreachable" 

486 

487 # lsstRGB returns an image in 0-1 scale it to the maximum value 

488 colorImage *= maxVal # type: ignore 

489 

490 # pack the joint mask back into a mask object 

491 lsstMask = Mask(width=jointMask.shape[1], height=jointMask.shape[0], planeDefs=maskDict) 

492 lsstMask.array = jointMask # type: ignore 

493 return Struct(outputRGB=colorImage.astype(dtype), outputRGBMask=lsstMask) # type: ignore 

494 

495 def runQuantum( 

496 self, 

497 butlerQC: QuantumContext, 

498 inputRefs: InputQuantizedConnection, 

499 outputRefs: OutputQuantizedConnection, 

500 ) -> None: 

501 imageRefs: list[DatasetRef] = inputRefs.inputCoadds 

502 sortedImages = self.makeInputsFromRefs(imageRefs, butlerQC) 

503 outputs = self.run(sortedImages) 

504 butlerQC.put(outputs, outputRefs) 

505 

506 def makeInputsFromRefs( 

507 self, refs: Iterable[DatasetRef], butler: Butler | QuantumContext 

508 ) -> dict[str, Exposure]: 

509 r"""Make valid inputs for the run method from butler references. 

510 

511 Parameters 

512 ---------- 

513 refs : `Iterable` of `DatasetRef` 

514 Some `Iterable` container of `Butler` `DatasetRef`\ s 

515 butler : `Butler` or `QuantumContext` 

516 This is the object that fetches the input data. 

517 

518 Returns 

519 ------- 

520 sortedImages : `dict` of `str` to `Exposure` 

521 A dictionary of `Exposure`\ s that keyed by the band they 

522 correspond to. 

523 """ 

524 sortedImages: dict[str, Exposure] = {} 

525 for ref in refs: 

526 key: str = cast(str, ref.dataId["band"]) 

527 image = butler.get(ref) 

528 sortedImages[key] = image 

529 return sortedImages 

530 

531 def makeInputsFromArrays(self, **kwargs) -> dict[int, DeferredDatasetHandle]: 

532 r"""Make valid inputs for the run method from numpy arrays. 

533 

534 Parameters 

535 ---------- 

536 kwargs : `NDArray` 

537 This is standard python kwargs where the left side of the equals 

538 is the data band, and the right side is the corresponding `NDArray` 

539 array. 

540 

541 Returns 

542 ------- 

543 sortedImages : `dict` of `str` to `Exposure` 

544 A dictionary of `Exposure`\ s that keyed by the band they 

545 correspond to. 

546 """ 

547 # ignore type because there aren't proper stubs for afw 

548 temp = {} 

549 for key, array in kwargs.items(): 

550 temp[key] = Exposure(Box2I(Point2I(0, 0), Extent2I(*array.shape)), dtype=array.dtype) 

551 temp[key].image.array[:] = array 

552 

553 return self.makeInputsFromExposures(**temp) 

554 

555 def makeInputsFromExposures(self, **kwargs) -> dict[int, DeferredDatasetHandle]: 

556 r"""Make valid inputs for the run method from `Exposure` objects. 

557 

558 Parameters 

559 ---------- 

560 kwargs : `Exposure` 

561 This is standard python kwargs where the left side of the equals 

562 is the data band, and the right side is the corresponding 

563 `Exposure`. 

564 

565 Returns 

566 ------- 

567 sortedImages : `dict` of `str` to `Exposure` 

568 A dictionary of `Exposure`\ s that keyed by the band they 

569 correspond to. 

570 """ 

571 sortedImages = {} 

572 for key, value in kwargs.items(): 

573 sortedImages[key] = value 

574 return sortedImages 

575 

576 

577class PrettyPictureBackgroundFixerConnections( 

578 PipelineTaskConnections, 

579 dimensions=("tract", "patch", "skymap", "band"), 

580 defaultTemplates={"coaddTypeName": "deep"}, 

581): 

582 inputCoadd = Input( 

583 doc=("Input coadd for which the background is to be removed"), 

584 name="{coaddTypeName}CoaddPsfMatched", 

585 storageClass="ExposureF", 

586 dimensions=("tract", "patch", "skymap", "band"), 

587 ) 

588 outputCoadd = Output( 

589 doc="The coadd with the background fixed and subtracted", 

590 name="pretty_picture_coadd_bg_subtracted", 

591 storageClass="ExposureF", 

592 dimensions=("tract", "patch", "skymap", "band"), 

593 ) 

594 

595 

596class PrettyPictureBackgroundFixerConfig( 

597 PipelineTaskConfig, pipelineConnections=PrettyPictureBackgroundFixerConnections 

598): 

599 pass 

600 

601 

602class PrettyPictureBackgroundFixerTask(PipelineTask): 

603 """Empirically flatten an images background. 

604 

605 Many astrophysical images have backgrounds with imperfections in them. 

606 This Task attempts to determine control points which are considered 

607 background values, and fits a radial basis function model to those 

608 points. This model is then subtracted off the image. 

609 

610 """ 

611 

612 _DefaultName = "prettyPictureBackgroundFixerTask" 

613 ConfigClass = PrettyPictureBackgroundFixerConfig 

614 

615 config: ConfigClass 

616 

617 def _neg_log_likelihood(self, params, x): 

618 """Calculate the negative log-likelihood for a Gaussian distribution. 

619 

620 This function computes the negative log-likelihood of a set of data `x` 

621 given a Gaussian distribution with parameters `mu` and `sigma`. It's 

622 designed to be used as the objective function for a minimization routine 

623 to find the best-fit Gaussian parameters. 

624 

625 Parameters 

626 ---------- 

627 params : `tuple` 

628 A tuple containing the mean (`mu`) and standard deviation (`sigma`) 

629 of the Gaussian distribution. 

630 x : `NDArray` 

631 The data samples for which to calculate the log-likelihood. 

632 

633 Returns 

634 ------- 

635 float 

636 The negative log-likelihood of the data given the Gaussian parameters. 

637 Returns infinity if sigma is non-positive or if the mean is less than 

638 the maximum value in x (to enforce the constraint that the Gaussian 

639 only models the lower tail of the distribution). 

640 """ 

641 mu, sigma = params 

642 if sigma <= 0: 

643 return np.inf 

644 M = np.max(x) 

645 if mu < M - 1e-8: # Allow for floating point precision issues 

646 return np.inf 

647 z = (x - mu) / sigma 

648 term = np.log(2) - np.log(sigma) + norm.logpdf(z) 

649 loglikelihood = np.sum(term) 

650 return -loglikelihood 

651 

652 def _tile_slices(self, arr, R, C): 

653 """Generate slices for tiling an array. 

654 

655 This function divides an array into a grid of tiles and returns a list of 

656 slice objects representing each tile. It handles cases where the array 

657 dimensions are not evenly divisible by the number of tiles in each 

658 dimension, distributing the remainder among the tiles. 

659 

660 Parameters 

661 ---------- 

662 arr : `NDArray` 

663 The input array to be tiled. Used only to determine the array's shape. 

664 R : `int` 

665 The number of tiles in the row dimension. 

666 C : `int` 

667 The number of tiles in the column dimension. 

668 

669 Returns 

670 ------- 

671 slices : `list` of `tuple` 

672 A list of tuples, where each tuple contains two `slice` objects 

673 representing the row and column slices for a single tile. 

674 """ 

675 M = arr.shape[0] 

676 N = arr.shape[1] 

677 

678 # Function to compute slices for a given dimension size and number of divisions 

679 def get_slices(total_size, num_divisions): 

680 base = total_size // num_divisions 

681 remainder = total_size % num_divisions 

682 slices = [] 

683 start = 0 

684 for i in range(num_divisions): 

685 end = start + base 

686 if i < remainder: 

687 end += 1 

688 slices.append((start, end)) 

689 start = end 

690 return slices 

691 

692 # Get row and column slices 

693 row_slices = get_slices(M, R) 

694 col_slices = get_slices(N, C) 

695 

696 # Generate all possible tile combinations of row and column slices 

697 tiles = [] 

698 for rs in row_slices: 

699 r_start, r_end = rs 

700 for cs in col_slices: 

701 c_start, c_end = cs 

702 tile_slice = (slice(r_start, r_end), slice(c_start, c_end)) 

703 tiles.append(tile_slice) 

704 

705 return tiles 

706 

707 def fixBackground(self, image): 

708 """Estimate and subtract the background from an image. 

709 

710 This function estimates the background level in an image using a median-based 

711 approach combined with Gaussian fitting and radial basis function interpolation. 

712 It aims to provide a more accurate background estimation than a simple median 

713 filter, especially in images with varying background levels. 

714 

715 Parameters 

716 ---------- 

717 image : `NDArray` 

718 The input image as a NumPy array. 

719 

720 Returns 

721 ------- 

722 numpy.ndarray 

723 An array representing the estimated background level across the image. 

724 """ 

725 # Find the median value in the image, which is likely to be 

726 # close to average background. Note this doesn't work well 

727 # in fields with high density or diffuse flux. 

728 maxLikely = np.median(image, axis=None) 

729 

730 # find all the pixels that are fainter than this 

731 # and find the std. This is just used as an initialization 

732 # parameter and doesn't need to be accurate. 

733 mask = image < maxLikely 

734 initial_std = (image[mask] - maxLikely).std() 

735 

736 # Don't do anything if there are no pixels to check 

737 if np.any(mask): 

738 # use a minimizer to determine best mu and sigma for a Gaussian 

739 # given only samples below the mean of the Gaussian. 

740 _, sigma_hat = halfnorm.fit(np.abs(image[mask] - maxLikely)) 

741 mu_hat = maxLikely 

742 else: 

743 mu_hat, sigma_hat = (maxLikely, 2 * initial_std) 

744 

745 # create a new masking threshold that is the determined 

746 # mean plus std from the fit 

747 threshhold = mu_hat + sigma_hat 

748 image_mask = image < threshhold 

749 

750 # create python slices that tile the image. 

751 tiles = self._tile_slices(image, 25, 25) 

752 

753 yloc = [] 

754 xloc = [] 

755 values = [] 

756 

757 # for each box find the middle position and the median background 

758 # value in the window. 

759 for xslice, yslice in tiles: 

760 ypos = (yslice.stop - yslice.start) / 2 + yslice.start 

761 xpos = (xslice.stop - xslice.start) / 2 + xslice.start 

762 yloc.append(ypos) 

763 xloc.append(xpos) 

764 window = image[yslice, xslice][image_mask[yslice, xslice]] 

765 if window.size > 0: 

766 value = np.median(window) 

767 else: 

768 value = 0 

769 values.append(value) 

770 

771 positions = np.meshgrid(np.arange(image.shape[0]), np.arange(image.shape[1])) 

772 # create an interpolant for the background and interpolate over the image. 

773 inter = RBFInterpolator( 

774 np.vstack((yloc, xloc)).T, values, kernel="thin_plate_spline", degree=4, smoothing=0.05 

775 ) 

776 backgrounds = inter(np.array(positions)[::-1].reshape(2, -1).T).reshape(image.shape) 

777 

778 return backgrounds 

779 

780 def run(self, inputCoadd: Exposure): 

781 """Estimate a background for an input Exposure and remove it. 

782 

783 Parameters 

784 ---------- 

785 inputCoadd : `Exposure` 

786 The exposure the background will be removed from. 

787 

788 Returns 

789 ------- 

790 result : `Struct` 

791 A `Struct` that contains the exposure with the background removed. 

792 This `Struct` will have an attribute named ``outputCoadd``. 

793 

794 """ 

795 background = self.fixBackground(inputCoadd.image.array) 

796 # create a copy to mutate 

797 output = ExposureF(inputCoadd, deep=True) 

798 output.image.array -= background 

799 return Struct(outputCoadd=output) 

800 

801 

802class PrettyPictureStarFixerConnections( 

803 PipelineTaskConnections, 

804 dimensions=("tract", "patch", "skymap"), 

805): 

806 inputCoadd = Input( 

807 doc=("Input coadd for which the background is to be removed"), 

808 name="pretty_picture_coadd_bg_subtracted", 

809 storageClass="ExposureF", 

810 dimensions=("tract", "patch", "skymap", "band"), 

811 multiple=True, 

812 ) 

813 outputCoadd = Output( 

814 doc="The coadd with the background fixed and subtracted", 

815 name="pretty_picture_coadd_fixed_stars", 

816 storageClass="ExposureF", 

817 dimensions=("tract", "patch", "skymap", "band"), 

818 multiple=True, 

819 ) 

820 

821 

822class PrettyPictureStarFixerConfig(PipelineTaskConfig, pipelineConnections=PrettyPictureStarFixerConnections): 

823 brightnessThresh = Field[float]( 

824 doc="The flux value below which pixels with SAT or NO_DATA bits will be ignored" 

825 ) 

826 

827 

828class PrettyPictureStarFixerTask(PipelineTask): 

829 """This class fixes up regions in an image where there is no, or bad data. 

830 

831 The fixes done by this task are overwhelmingly comprised of the cores of 

832 bright stars for which there is no data. 

833 """ 

834 

835 _DefaultName = "prettyPictureStarFixerTask" 

836 ConfigClass = PrettyPictureStarFixerConfig 

837 

838 config: ConfigClass 

839 

840 def run(self, inputs: Mapping[str, ExposureF]) -> Struct: 

841 """Fix areas in an image where this is no data, most likely to be 

842 the cores of bright stars. 

843 

844 Because we want to have consistent fixes accross bands, this method 

845 relies on supplying all bands and fixing pixels that are marked 

846 as having a defect in any band even if within one band there is 

847 no issue. 

848 

849 Parameters 

850 ---------- 

851 inputs : `Mapping` of `str` to `ExposureF` 

852 This mapping has keys of band as a `str` and the corresponding 

853 ExposureF as a value. 

854 

855 Returns 

856 ------- 

857 results : `Struct` of `Mapping` of `str` to `ExposureF` 

858 A `Struct` that has a mapping of band to `ExposureF`. The `Struct` 

859 has an attribute named ``results``. 

860 

861 """ 

862 # make the joint mask of all the channels 

863 doJointMaskInit = True 

864 for imageExposure in inputs.values(): 

865 maskDict = imageExposure.mask.getMaskPlaneDict() 

866 if doJointMaskInit: 

867 jointMask = np.zeros(imageExposure.mask.array.shape, dtype=imageExposure.mask.array.dtype) 

868 doJointMaskInit = False 

869 jointMask |= imageExposure.mask.array 

870 

871 sat_bit = maskDict["SAT"] 

872 no_data_bit = maskDict["NO_DATA"] 

873 together = (jointMask & 2**sat_bit).astype(bool) | (jointMask & 2**no_data_bit).astype(bool) 

874 

875 # use the last imageExposure as it is likely close enough across all bands 

876 bright_mask = imageExposure.image.array > self.config.brightnessThresh 

877 

878 # dilate the mask a bit, this helps get a bit fainter mask without starting 

879 # to include pixels in an irregular shape, as only the star cores should be 

880 # fixed. 

881 both = together & bright_mask 

882 struct = np.array(((0, 1, 0), (1, 1, 1), (0, 1, 0)), dtype=bool) 

883 both = binary_dilation(both, struct, iterations=4).astype(bool) 

884 

885 # do the actual fixing of values 

886 results = {} 

887 for band, imageExposure in inputs.items(): 

888 if np.sum(both) > 0: 

889 inpainted = inpaint_biharmonic(imageExposure.image.array, both, split_into_regions=True) 

890 imageExposure.image.array[both] = inpainted[both] 

891 results[band] = imageExposure 

892 return Struct(results=results) 

893 

894 def runQuantum(