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

371 statements  

« prev     ^ index     » next       coverage.py v7.13.1, created at 2026-01-07 01:55 -0800

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 "PrettyPictureTask", 

26 "PrettyPictureConnections", 

27 "PrettyPictureConfig", 

28 "PrettyMosaicTask", 

29 "PrettyMosaicConnections", 

30 "PrettyMosaicConfig", 

31 "PrettyPictureBackgroundFixerConfig", 

32 "PrettyPictureBackgroundFixerTask", 

33 "PrettyPictureStarFixerConfig", 

34 "PrettyPictureStarFixerTask", 

35) 

36 

37from collections.abc import Iterable, Mapping 

38from lsst.afw.image import ExposureF 

39import numpy as np 

40from typing import TYPE_CHECKING, cast, Any 

41from lsst.skymap import BaseSkyMap 

42 

43from scipy.stats import norm 

44from scipy.ndimage import binary_dilation 

45from scipy.optimize import minimize 

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="{coaddTypeName}CoaddPsfMatched", 

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 

256 def setDefaults(self): 

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

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

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

260 return super().setDefaults() 

261 

262 

263class PrettyPictureTask(PipelineTask): 

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

265 

266 _DefaultName = "prettyPictureTask" 

267 ConfigClass = PrettyPictureConfig 

268 

269 config: ConfigClass 

270 

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

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

273 

274 Parameters 

275 ---------- 

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

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

278 

279 Returns 

280 ------- 

281 result : `Struct` 

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

283 RGB image construction. The struct will have the attributes 

284 outputRGBImage and outputRGBMask. Each of the outputs will 

285 be a `NDarray` object. 

286 

287 Notes 

288 ----- 

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

290 makeInputsFrom* methods. 

291 """ 

292 channels = {} 

293 shape = (0, 0) 

294 jointMask: None | NDArray = None 

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

296 doJointMaskInit = False 

297 if jointMask is None: 

298 doJointMask = True 

299 doJointMaskInit = True 

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

301 imageArray = imageExposure.image.array 

302 # run all the plugins designed for array based interaction 

303 for plug in plugins.channel(): 

304 imageArray = plug( 

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

306 ).astype(np.float32) 

307 channels[channel] = imageArray 

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

309 # happen in each loop. 

310 shape = imageArray.shape 

311 maskDict = imageExposure.mask.getMaskPlaneDict() 

312 if doJointMaskInit: 

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

314 doJointMaskInit = False 

315 if doJointMask: 

316 jointMask |= imageExposure.mask.array 

317 

318 # mix the images to RGB 

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

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

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

322 

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

324 mix = self.config.channelConfig[band] 

325 if mix.r: 

326 imageRArray += mix.r * image 

327 if mix.g: 

328 imageGArray += mix.g * image 

329 if mix.b: 

330 imageBArray += mix.b * image 

331 

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

333 box: Box2I = exposure.getBBox() 

334 boxCenter = box.getCenter() 

335 try: 

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

337 except Exception: 

338 psf = None 

339 

340 # assert for typing reasons 

341 assert jointMask is not None 

342 # Run any image level correction plugins 

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

344 colorImage[:, :, 0] = imageRArray 

345 colorImage[:, :, 1] = imageGArray 

346 colorImage[:, :, 2] = imageBArray 

347 for plug in plugins.partial(): 

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

349 

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

351 # and not typed. 

352 colorImage = lsstRGB( 

353 colorImage[:, :, 0], 

354 colorImage[:, :, 1], 

355 colorImage[:, :, 2], 

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

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

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

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

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

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

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

363 doRemapGamut=self.config.doRemapGamut, 

364 gamutMethod=self.config.gamutMethod, 

365 ) 

366 

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

368 maxVal: int | float 

369 match self.config.arrayType: 

370 case "uint8": 

371 dtype = np.uint8 

372 maxVal = 255 

373 case "uint16": 

374 dtype = np.uint16 

375 maxVal = 65535 

376 case "half": 

377 dtype = np.half 

378 maxVal = 1.0 

379 case "float": 

380 dtype = np.float32 

381 maxVal = 1.0 

382 case _: 

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

384 

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

386 colorImage *= maxVal # type: ignore 

387 

388 # pack the joint mask back into a mask object 

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

390 lsstMask.array = jointMask # type: ignore 

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

392 

393 def runQuantum( 

394 self, 

395 butlerQC: QuantumContext, 

396 inputRefs: InputQuantizedConnection, 

397 outputRefs: OutputQuantizedConnection, 

398 ) -> None: 

399 imageRefs: list[DatasetRef] = inputRefs.inputCoadds 

400 sortedImages = self.makeInputsFromRefs(imageRefs, butlerQC) 

401 outputs = self.run(sortedImages) 

402 butlerQC.put(outputs, outputRefs) 

403 

404 def makeInputsFromRefs( 

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

406 ) -> dict[str, Exposure]: 

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

408 

409 Parameters 

410 ---------- 

411 refs : `Iterable` of `DatasetRef` 

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

413 butler : `Butler` or `QuantumContext` 

414 This is the object that fetches the input data. 

415 

416 Returns 

417 ------- 

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

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

420 correspond to. 

421 """ 

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

423 for ref in refs: 

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

425 image = butler.get(ref) 

426 sortedImages[key] = image 

427 return sortedImages 

428 

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

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

431 

432 Parameters 

433 ---------- 

434 kwargs : `NDArray` 

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

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

437 array. 

438 

439 Returns 

440 ------- 

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

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

443 correspond to. 

444 """ 

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

446 temp = {} 

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

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

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

450 

451 return self.makeInputsFromExposures(**temp) 

452 

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

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

455 

456 Parameters 

457 ---------- 

458 kwargs : `Exposure` 

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

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

461 `Exposure`. 

462 

463 Returns 

464 ------- 

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

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

467 correspond to. 

468 """ 

469 sortedImages = {} 

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

471 sortedImages[key] = value 

472 return sortedImages 

473 

474 

475class PrettyPictureBackgroundFixerConnections( 

476 PipelineTaskConnections, 

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

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

479): 

480 inputCoadd = Input( 

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

482 name="{coaddTypeName}CoaddPsfMatched", 

483 storageClass="ExposureF", 

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

485 ) 

486 outputCoadd = Output( 

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

488 name="pretty_picture_coadd_bg_subtracted", 

489 storageClass="ExposureF", 

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

491 ) 

492 

493 

494class PrettyPictureBackgroundFixerConfig( 

495 PipelineTaskConfig, pipelineConnections=PrettyPictureBackgroundFixerConnections 

496): 

497 pass 

498 

499 

500class PrettyPictureBackgroundFixerTask(PipelineTask): 

501 """Empirically flatten an images background. 

502 

503 Many astrophysical images have backgrounds with imperfections in them. 

504 This Task attempts to determine control points which are considered 

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

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

507 

508 """ 

509 

510 _DefaultName = "prettyPictureBackgroundFixerTask" 

511 ConfigClass = PrettyPictureBackgroundFixerConfig 

512 

513 config: ConfigClass 

514 

515 def _neg_log_likelihood(self, params, x): 

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

517 

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

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

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

521 to find the best-fit Gaussian parameters. 

522 

523 Parameters 

524 ---------- 

525 params : `tuple` 

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

527 of the Gaussian distribution. 

528 x : `NDArray` 

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

530 

531 Returns 

532 ------- 

533 float 

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

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

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

537 only models the lower tail of the distribution). 

538 """ 

539 mu, sigma = params 

540 if sigma <= 0: 

541 return np.inf 

542 M = np.max(x) 

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

544 return np.inf 

545 z = (x - mu) / sigma 

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

547 loglikelihood = np.sum(term) 

548 return -loglikelihood 

549 

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

551 """Generate slices for tiling an array. 

552 

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

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

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

556 dimension, distributing the remainder among the tiles. 

557 

558 Parameters 

559 ---------- 

560 arr : `NDArray` 

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

562 R : `int` 

563 The number of tiles in the row dimension. 

564 C : `int` 

565 The number of tiles in the column dimension. 

566 

567 Returns 

568 ------- 

569 slices : `list` of `tuple` 

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

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

572 """ 

573 M = arr.shape[0] 

574 N = arr.shape[1] 

575 

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

577 def get_slices(total_size, num_divisions): 

578 base = total_size // num_divisions 

579 remainder = total_size % num_divisions 

580 slices = [] 

581 start = 0 

582 for i in range(num_divisions): 

583 end = start + base 

584 if i < remainder: 

585 end += 1 

586 slices.append((start, end)) 

587 start = end 

588 return slices 

589 

590 # Get row and column slices 

591 row_slices = get_slices(M, R) 

592 col_slices = get_slices(N, C) 

593 

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

595 tiles = [] 

596 for rs in row_slices: 

597 r_start, r_end = rs 

598 for cs in col_slices: 

599 c_start, c_end = cs 

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

601 tiles.append(tile_slice) 

602 

603 return tiles 

604 

605 def fixBackground(self, image): 

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

607 

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

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

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

611 filter, especially in images with varying background levels. 

612 

613 Parameters 

614 ---------- 

615 image : `NDArray` 

616 The input image as a NumPy array. 

617 

618 Returns 

619 ------- 

620 numpy.ndarray 

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

622 """ 

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

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

625 # in fields with high density or diffuse flux. 

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

627 

628 # find all the pixels that are fainter than this 

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

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

631 mask = image < maxLikely 

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

633 

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

635 if np.any(mask): 

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

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

638 result = minimize( 

639 self._neg_log_likelihood, 

640 (maxLikely, initial_std), 

641 args=(image[mask]), 

642 bounds=((maxLikely, None), (1e-8, None)), 

643 ) 

644 mu_hat, sigma_hat = result.x 

645 else: 

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

647 

648 # create a new masking threshold that is the determined 

649 # mean plus std from the fit 

650 threshhold = mu_hat + sigma_hat 

651 image_mask = image < threshhold 

652 

653 # create python slices that tile the image. 

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

655 

656 yloc = [] 

657 xloc = [] 

658 values = [] 

659 

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

661 # value in the window. 

662 for xslice, yslice in tiles: 

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

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

665 yloc.append(ypos) 

666 xloc.append(xpos) 

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

668 if window.size > 0: 

669 value = np.median(window) 

670 else: 

671 value = 0 

672 values.append(value) 

673 

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

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

676 inter = RBFInterpolator( 

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

678 ) 

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

680 

681 return backgrounds 

682 

683 def run(self, inputCoadd: Exposure): 

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

685 

686 Parameters 

687 ---------- 

688 inputCoadd : `Exposure` 

689 The exposure the background will be removed from. 

690 

691 Returns 

692 ------- 

693 result : `Struct` 

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

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

696 

697 """ 

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

699 # create a copy to mutate 

700 output = ExposureF(inputCoadd, deep=True) 

701 output.image.array -= background 

702 return Struct(outputCoadd=output) 

703 

704 

705class PrettyPictureStarFixerConnections( 

706 PipelineTaskConnections, 

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

708): 

709 inputCoadd = Input( 

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

711 name="pretty_picture_coadd_bg_subtracted", 

712 storageClass="ExposureF", 

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

714 multiple=True, 

715 ) 

716 outputCoadd = Output( 

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

718 name="pretty_picture_coadd_fixed_stars", 

719 storageClass="ExposureF", 

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

721 multiple=True, 

722 ) 

723 

724 

725class PrettyPictureStarFixerConfig(PipelineTaskConfig, pipelineConnections=PrettyPictureStarFixerConnections): 

726 brightnessThresh = Field[float]( 

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

728 ) 

729 

730 

731class PrettyPictureStarFixerTask(PipelineTask): 

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

733 

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

735 bright stars for which there is no data. 

736 """ 

737 

738 _DefaultName = "prettyPictureStarFixerTask" 

739 ConfigClass = PrettyPictureStarFixerConfig 

740 

741 config: ConfigClass 

742 

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

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

745 the cores of bright stars. 

746 

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

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

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

750 no issue. 

751 

752 Parameters 

753 ---------- 

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

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

756 ExposureF as a value. 

757 

758 Returns 

759 ------- 

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

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

762 has an attribute named ``results``. 

763 

764 """ 

765 # make the joint mask of all the channels 

766 doJointMaskInit = True 

767 for imageExposure in inputs.values(): 

768 maskDict = imageExposure.mask.getMaskPlaneDict() 

769 if doJointMaskInit: 

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

771 doJointMaskInit = False 

772 jointMask |= imageExposure.mask.array 

773 

774 sat_bit = maskDict["SAT"] 

775 no_data_bit = maskDict["NO_DATA"] 

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

777 

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

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

780 

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

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

783 # fixed. 

784 both = together & bright_mask 

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

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

787 

788 # do the actual fixing of values 

789 results = {} 

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

791 if np.sum(both) > 0: 

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

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

794 results[band] = imageExposure 

795 return Struct(results=results) 

796 

797 def runQuantum( 

798 self, 

799 butlerQC: QuantumContext, 

800 inputRefs: InputQuantizedConnection, 

801 outputRefs: OutputQuantizedConnection, 

802 ) -> None: 

803 refs = inputRefs.inputCoadd 

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

805 for ref in refs: 

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

807 image = butlerQC.get(ref) 

808 sortedImages[key] = image 

809 

810 outputs = self.run(sortedImages).results 

811 sortedOutputs = {} 

812 for ref in outputRefs.outputCoadd: 

813 sortedOutputs[ref.dataId["band"]] = ref 

814 

815 for band, data in outputs.items(): 

816 butlerQC.put(data, sortedOutputs[band]) 

817 

818 

819class PrettyMosaicConnections(PipelineTaskConnections, dimensions=("tract", "skymap")): 

820 inputRGB = Input( 

821 doc="Individual RGB images that are to go into the mosaic", 

822 name="rgb_picture_array", 

823 storageClass="NumpyArray", 

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

825 multiple=True, 

826 deferLoad=True, 

827 ) 

828 

829 skyMap = Input( 

830 doc="The skymap which the data has been mapped onto", 

831 storageClass="SkyMap", 

832 name=BaseSkyMap.SKYMAP_DATASET_TYPE_NAME, 

833 dimensions=("skymap",), 

834 ) 

835 

836 inputRGBMask = Input( 

837 doc="Individual RGB images that are to go into the mosaic", 

838 name="rgb_picture_mask", 

839 storageClass="Mask", 

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

841 multiple=True, 

842 deferLoad=True, 

843 ) 

844 

845 outputRGBMosaic = Output( 

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

847 name="rgb_mosaic_array", 

848 storageClass="NumpyArray", 

849 dimensions=("tract", "skymap"), 

850 ) 

851 

852 

853class PrettyMosaicConfig(PipelineTaskConfig, pipelineConnections=PrettyMosaicConnections): 

854 binFactor = Field[int](doc="The factor to bin by when producing the mosaic") 

855 

856 

857class PrettyMosaicTask(PipelineTask): 

858 """Combines multiple RGB arrays into one mosaic.""" 

859 

860 _DefaultName = "prettyMosaicTask" 

861 ConfigClass = PrettyMosaicConfig 

862 

863 config: ConfigClass 

864 

865 def run( 

866 self, 

867 inputRGB: Iterable[DeferredDatasetHandle], 

868 skyMap: BaseSkyMap, 

869 inputRGBMask: Iterable[DeferredDatasetHandle], 

870 ) -> Struct: 

871 r"""Assemble individual `NDArrays` into a mosaic. 

872 

873 Each input is a `DeferredDatasetHandle` because they're loaded in one 

874 at a time to be placed into the mosaic to save memory. 

875 

876 Parameters 

877 ---------- 

878 inputRGB : `Iterable` of `DeferredDatasetHandle` 

879 `DeferredDatasetHandle`\ s pointing to RGB `NDArrays`. 

880 skyMap : `BaseSkyMap` 

881 The skymap that defines the relative position of each of the input 

882 images. 

883 inputRGBMask : `Iterable` of `DeferredDatasetHandle` 

884 `DeferredDatasetHandle`\ s pointing to masks for each of the 

885 corresponding images. 

886 

887 Returns 

888 ------- 

889 result : `Struct` 

890 The `Struct` containing the combined mosaic. The `Struct` has 

891 and attribute named ``outputRGBMosaic``. 

892 """ 

893 # create the bounding region 

894 newBox = Box2I() 

895 # store the bounds as they are retrieved from the skymap 

896 boxes = [] 

897 tractMaps = [] 

898 for handle in inputRGB: 

899 dataId = handle.dataId 

900 tractInfo: TractInfo = skyMap[dataId["tract"]] 

901 patchInfo: PatchInfo = tractInfo[dataId["patch"]] 

902 bbox = patchInfo.getOuterBBox() 

903 boxes.append(bbox) 

904 newBox.include(bbox) 

905 tractMaps.append(tractInfo) 

906 

907 # fixup the boxes to be smaller if needed, and put the origin at zero, 

908 # this must be done after constructing the complete outer box 

909 modifiedBoxes = [] 

910 origin = newBox.getBegin() 

911 for iterBox in boxes: 

912 localOrigin = iterBox.getBegin() - origin 

913 localOrigin = Point2I( 

914 x=int(np.floor(localOrigin.x / self.config.binFactor)), 

915 y=int(np.floor(localOrigin.y / self.config.binFactor)), 

916 ) 

917 localExtent = Extent2I( 

918 x=int(np.floor(iterBox.getWidth() / self.config.binFactor)), 

919 y=int(np.floor(iterBox.getHeight() / self.config.binFactor)), 

920 ) 

921 tmpBox = Box2I(localOrigin, localExtent) 

922 modifiedBoxes.append(tmpBox) 

923 boxes = modifiedBoxes 

924 

925 # scale the container box 

926 newBoxOrigin = Point2I(0, 0) 

927 newBoxExtent = Extent2I( 

928 x=int(np.floor(newBox.getWidth() / self.config.binFactor)), 

929 y=int(np.floor(newBox.getHeight() / self.config.binFactor)), 

930 ) 

931 newBox = Box2I(newBoxOrigin, newBoxExtent) 

932 

933 # Allocate storage for the mosaic 

934 self.imageHandle = tempfile.NamedTemporaryFile() 

935 self.maskHandle = tempfile.NamedTemporaryFile() 

936 consolidatedImage = None 

937 consolidatedMask = None 

938 

939 # Actually assemble the mosaic 

940 maskDict = {} 

941 tmpImg = None 

942 for box, handle, handleMask, tractInfo in zip(boxes, inputRGB, inputRGBMask, tractMaps): 

943 rgb = handle.get() 

944 rgbMask = handleMask.get() 

945 maskDict = rgbMask.getMaskPlaneDict() 

946 # allocate the memory for the mosaic 

947 if consolidatedImage is None: 

948 consolidatedImage = np.memmap( 

949 self.imageHandle.name, 

950 mode="w+", 

951 shape=(newBox.getHeight(), newBox.getWidth(), 3), 

952 dtype=rgb.dtype, 

953 ) 

954 if consolidatedMask is None: 

955 consolidatedMask = np.memmap( 

956 self.maskHandle.name, 

957 mode="w+", 

958 shape=(newBox.getHeight(), newBox.getWidth()), 

959 dtype=rgbMask.array.dtype, 

960 ) 

961 

962 if self.config.binFactor > 1: 

963 # opencv wants things in x, y dimensions 

964 shape = tuple(box.getDimensions())[::-1] 

965 rgb = cv2.resize( 

966 rgb, 

967 dst=None, 

968 dsize=shape, 

969 fx=shape[0] / self.config.binFactor, 

970 fy=shape[1] / self.config.binFactor, 

971 ) 

972 rgbMask = cv2.resize( 

973 rgbMask.array.astype(np.float32), 

974 dst=None, 

975 dsize=shape, 

976 fx=shape[0] / self.config.binFactor, 

977 fy=shape[1] / self.config.binFactor, 

978 ) 

979 existing = ~np.all(consolidatedImage[*box.slices] == 0, axis=2) 

980 if tmpImg is None or tmpImg.shape != rgb.shape: 

981 ramp = np.linspace(0, 1, tractInfo.patch_border * 2) 

982 tmpImg = np.zeros(rgb.shape[:2]) 

983 tmpImg[: tractInfo.patch_border * 2, :] = np.repeat( 

984 np.expand_dims(ramp, 1), tmpImg.shape[1], axis=1 

985 ) 

986 

987 tmpImg[-1 * tractInfo.patch_border * 2 :, :] = np.repeat( # noqa: E203 

988 np.expand_dims(1 - ramp, 1), tmpImg.shape[1], axis=1 

989 ) 

990 tmpImg[:, : tractInfo.patch_border * 2] = np.repeat( 

991 np.expand_dims(ramp, 0), tmpImg.shape[0], axis=0 

992 ) 

993 

994 tmpImg[:, -1 * tractInfo.patch_border * 2 :] = np.repeat( # noqa: E203 

995 np.expand_dims(1 - ramp, 0), tmpImg.shape[0], axis=0 

996 ) 

997 tmpImg = np.repeat(np.expand_dims(tmpImg, 2), 3, axis=2) 

998 

999 consolidatedImage[*box.slices][~existing, :] = rgb[~existing, :] 

1000 consolidatedImage[*box.slices][existing, :] = ( 

1001 tmpImg[existing] * rgb[existing] 

1002 + (1 - tmpImg[existing]) * consolidatedImage[*box.slices][existing, :] 

1003 ) 

1004 

1005 tmpMask = np.zeros_like(rgbMask.array) 

1006 tmpMask[existing] = np.bitwise_or( 

1007 rgbMask.array[existing], consolidatedMask[*box.slices][existing] 

1008 ) 

1009 tmpMask[~existing] = rgbMask.array[~existing] 

1010 consolidatedMask[*box.slices] = tmpMask 

1011 

1012 for plugin in plugins.full(): 

1013 if consolidatedImage is not None and consolidatedMask is not None: 

1014 consolidatedImage = plugin(consolidatedImage, consolidatedMask, maskDict) 

1015 # If consolidated image still None, that means there was no work to do. 

1016 # Return an empty image instead of letting this task fail. 

1017 if consolidatedImage is None: 

1018 consolidatedImage = np.zeros((0, 0, 0), dtype=np.uint8) 

1019 

1020 return Struct(outputRGBMosaic=consolidatedImage) 

1021 

1022 def runQuantum( 

1023 self, 

1024 butlerQC: QuantumContext, 

1025 inputRefs: InputQuantizedConnection, 

1026 outputRefs: OutputQuantizedConnection, 

1027 ) -> None: 

1028 inputs = butlerQC.get(inputRefs) 

1029 outputs = self.run(**inputs) 

1030 butlerQC.put(outputs, outputRefs) 

1031 if hasattr(self, "imageHandle"): 

1032 self.imageHandle.close() 

1033 if hasattr(self, "maskHandle"): 

1034 self.maskHandle.close() 

1035 

1036 def makeInputsFromArrays( 

1037 self, inputs: Iterable[tuple[Mapping[str, Any], NDArray]] 

1038 ) -> Iterable[DeferredDatasetHandle]: 

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

1040 

1041 Parameters 

1042 ---------- 

1043 inputs : `Iterable` of `tuple` of `Mapping` and `NDArray` 

1044 An iterable where each element is a tuble with the first 

1045 element is a mapping that corresponds to an arrays dataId, 

1046 and the second is an `NDArray`. 

1047 

1048 Returns 

1049 ------- 

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

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

1052 correspond to. 

1053 """ 

1054 structuredInputs = [] 

1055 for dataId, array in inputs: 

1056 structuredInputs.append(InMemoryDatasetHandle(inMemoryDataset=array, **dataId)) 

1057 

1058 return structuredInputs