Coverage for python / lsst / pipe / tasks / rgb2hips / _high_order_hips.py: 0%
151 statements
« prev ^ index » next coverage.py v7.13.4, created at 2026-03-05 23:19 -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/>.
21from __future__ import annotations
23__all__ = ("HighOrderHipsTaskConnections", "HighOrderHipsTaskConfig", "HighOrderHipsTask")
25import numpy as np
26from enum import Enum
27from numpy.typing import NDArray
29from lsst.afw.geom import makeHpxWcs
30from lsst.pipe.base import (
31 PipelineTask,
32 PipelineTaskConfig,
33 PipelineTaskConnections,
34 Struct,
35 QuantumContext,
36 InputQuantizedConnection,
37 OutputQuantizedConnection,
38)
39from lsst.pex.config import ConfigField, Field, ChoiceField
40from lsst.pipe.base.connectionTypes import Input, Output
41from lsst.skymap import BaseSkyMap
42from lsst.afw.geom import SkyWcs
43from lsst.geom import Box2I, Point2I, Extent2I
44from lsst.afw.math import Warper
45from lsst.daf.butler import DeferredDatasetHandle
46from lsst.afw.image import ImageF
47from lsst.resources import ResourcePath
49from collections.abc import Iterable
50from lsst.sphgeom import RangeSet
52import cv2
54from ._utils import _write_hips_image
57class ColorChannel(Enum):
58 """Enum representing color channels for image processing."""
60 RED = 0
61 GREEN = 1
62 BLUE = 2
65class HighOrderHipsTaskConnections(PipelineTaskConnections, dimensions=("healpix8",)):
66 input_images = Input(
67 doc="Color images which are to be turned into hips tiles",
68 name="rgb_picture_array",
69 storageClass="NumpyArray",
70 dimensions=("tract", "patch", "skymap"),
71 multiple=True,
72 deferLoad=True,
73 )
74 skymap = Input(
75 doc="The skymap which the data has been mapped onto",
76 storageClass="SkyMap",
77 name=BaseSkyMap.SKYMAP_DATASET_TYPE_NAME,
78 dimensions=("skymap",),
79 )
80 output_hpx = Output(
81 doc="Healpix tiles at order 8, but binned to 256x256",
82 name="rgb_picture_hips8",
83 storageClass="NumpyArray",
84 dimensions=("healpix8",),
85 )
88class HighOrderHipsTaskConfig(PipelineTaskConfig, pipelineConnections=HighOrderHipsTaskConnections):
89 """Configuration class for the HighOrderHipsTask pipeline task."""
91 hips_order = 8
92 """HealPix order to generate tiles for."""
93 warp = ConfigField[Warper.ConfigClass](
94 doc="Warper configuration",
95 )
96 hips_base_uri = Field[str](
97 doc="URI to HiPS base for output.",
98 optional=False,
99 )
100 color_ordering = Field[str](
101 doc=(
102 "A string of the astrophysical bands that correspond to the RGB channels in the color image "
103 "inputs to high_order_hips task. This is in making the hips metadata"
104 ),
105 optional=False,
106 )
107 file_extension = ChoiceField[str](
108 doc="Extension for the presisted image.",
109 allowed={"png": "Use the png image extension", "webp": "Use the webp image extension"},
110 default="png",
111 )
112 array_type = ChoiceField[str](
113 doc="The dataset type for the output image array",
114 default="uint8",
115 allowed={
116 "uint8": "Use 8 bit arrays, 255 max",
117 "uint16": "Use 16 bit arrays, 65535 max",
118 "half": "Use 16 bit float arrays, 1 max",
119 "float": "Use 32 bit float arrays, 1 max",
120 },
121 )
123 def setDefaults(self):
124 self.warp.warpingKernelName = "lanczos5"
127class HighOrderHipsTask(PipelineTask):
128 """Pipeline task that generates high-order HealPix tiles from RGB images.
130 Of Note; This task has special dispensation to write "out-of-tree" to a
131 location not within the butler. DO NOT model other tasks on this one.
133 This task takes in RGB images generated on a tract patch grid. It assembles
134 them into a 4096 x 4096 image aligned with the wcs coordinates of hips
135 order 8 pixels. This is then divided up into an 8x8 grid to produce 512x512
136 images at hips order 11. The images is then resampled using lanczos order 4
137 such that the image is half the size. The original image is then divided
138 into a 4x4 grid to produce hips images at order 10. The process is repeated
139 to produce hips images at order 9, and finally the image is resampled down
140 to 512x512 and saved out at hips order 8.
142 The order 8 image is resampled one more time to 256x256 and presisted by
143 the butler for later consumption in the `LowOrderHipsTask`.
145 The difference at producding wcs at order 8 and working up to 11, is tested
146 to be less than 6 decimal places when converting ra dec to pixel coordinates,
147 and even that is likely to be due to differences in warping kernels,
148 and not an intrinsic error. Doing processing like this allows hips generation
149 to be more effectively split across compute nodes.
150 """
152 _DefaultName = "highOrderHipsTask"
153 ConfigClass = HighOrderHipsTaskConfig
155 config: ConfigClass
157 def __init__(self, **kwargs):
158 super().__init__(**kwargs)
159 self.warper = Warper.fromConfig(self.config.warp)
161 # Set the base resource path that will be used for all outputs
162 self.hips_base_path = ResourcePath(self.config.hips_base_uri, forceDirectory=True)
163 self.hips_base_path = self.hips_base_path.join(
164 f"color_{self.config.color_ordering}", forceDirectory=True
165 )
167 def run(self, input_images: Iterable[tuple[NDArray, SkyWcs, Box2I]], healpix_id) -> Struct:
168 """Main execution method for generating HealPix tiles.
170 Parameters
171 ----------
172 input_images : Iterable[tuple[NDArray, SkyWcs, Box2I]]
173 Iterable of tuples containing image data, WCS, and bounding box information.
174 healpix_id : int
175 The HealPix order 8 ID to process.
177 Returns
178 -------
179 Struct
180 Output structure containing the processed HealPix order 8 tile.
181 This has been downsampled to 256x256 corresponding to a quarter of a healpix
182 order 7 image.
183 """
184 # Make the WCS for the transform, intentionally over-sampled to shift order 12.
185 # This creates as 4096 x 4096 image that can be broken apart to form the higher
186 # orders, binning each as needed
187 target_wcs = makeHpxWcs(8, healpix_id, 12, False)
189 # construct a bounding box that holds the warping results for each channel
190 exp_bbox = Box2I(corner=Point2I(0, 0), dimensions=Extent2I(2**12, 2**12))
192 output_array_hpx = np.zeros((4096, 4096, 3), dtype=np.float32)
193 output_array_hpx[:, :, :] = np.nan
195 self.log.info("Warping input exposures and populating hpx8 super tile.")
196 # Need to loop over input arrays then channel
197 # Warp and combine input images into the HealPix tile
198 for input_image, in_wcs, in_box in input_images:
199 tmp_image = ImageF(in_box)
200 in_image: NDArray = input_image
202 # Normalize image data based on dtype
203 match in_image.dtype:
204 case np.uint8:
205 in_image = in_image.astype(np.float32) / 255.0
206 case np.uint16:
207 in_image = in_image.astype(np.float32) / 65535
208 case np.float16:
209 in_image = in_image.astype(np.float32)
211 # Process each color channel separately
212 for channel in ColorChannel:
213 # existing data
214 existing = output_array_hpx[..., channel.value]
216 # construct an Exposure object from one channel in the array
217 channel_array = in_image[..., channel.value]
218 tmp_image.array[:, :] = channel_array
220 # Warp the image to the target WCS
221 warpped = self.warper.warpImage(target_wcs, tmp_image, in_wcs, maxBBox=exp_bbox)
222 warpped_box_slices = warpped.getBBox().slices
224 # Update the output array with valid (non-NaN) values
225 are_warpped = np.isfinite(warpped.array)
226 existing[warpped_box_slices][are_warpped] = warpped.array[are_warpped]
228 # Replace any remaining NaN values with zeros
229 output_array_hpx[np.isnan(output_array_hpx)] = 0
231 # Flip the y-axis to match HealPix indexing
232 output_array_hpx = output_array_hpx[::-1, :, :]
234 # Generate tiles for different HealPix orders using Lanczos resampling instead of binning.
235 # This handles how intensities should change as the hips level changes.
236 #
237 # what this does is take a single 4096 x 4096 image and resamples it in a courser grain such
238 # that the output pixels correspond to a 4x4 grid of hips pixels at an increasingly lower scale.
239 # This works because hips is a hierarchy of tiles all contained in the same area of the sky.
240 # This allows us to generate all the output images by resampling the inputs and saves the time
241 # required to generate whole new images at each scale.
242 #
243 # The loop variables are the resampling factor, the hips order, and the number of sub-divisions
244 # a pixel has gone through (used to determine quadrant).
245 for zoom, hips_level, factor in zip((0, 2, 4, 8), (11, 10, 9, 8), (3, 2, 1, 0)):
246 self.log.info("Generating tiles for hxp level %d", hips_level)
247 if zoom:
248 size = 4096 // zoom
249 binned_array = cv2.resize(output_array_hpx, (size, size), interpolation=cv2.INTER_LANCZOS4)
250 else:
251 binned_array = output_array_hpx
252 # always create blocks of 512x512 as that is native shift order 9 size
253 #
254 # Figure out the hips pixel ids at this hips order. This is complicated because each hipx pixel
255 # turns into 4 at a higher level, but must be in a specific order to correspond to how the data
256 # is layed out in an y,x grid. So if a hips order 8 pixel A turns into four pixels b,c,d,e, they
257 # are layed out like [[b,d], [c,e]]. This is true for every pixel as you go up in order. So
258 # if you start at order 8 with one pixel, you need to do order 9 and calculate the layout. Then
259 # for each order 9 pixel, do the same to get the layout in order 10, etc. This leaves a grid
260 # of pixels that are the ids of the corresponding 512,512 sub grid pixel in the input image.
261 tmp_pixels = np.array([[healpix_id]])
262 for _ in range(factor):
263 tmp_array = np.zeros(np.array(tmp_pixels.shape) * 2)
264 for ii in range(tmp_pixels.shape[0]):
265 for jj in range(tmp_pixels.shape[1]):
266 tmp_array_view = tmp_array[ii * 2 : ii * 2 + 2, jj * 2 : jj * 2 + 2]
267 tmp_range_set = RangeSet(int(tmp_pixels[ii, jj]))
268 tmp_array_view[:, :] = (
269 np.array([x for x in range(*tmp_range_set.scaled(4)[0])], dtype=int)[[0, 2, 1, 3]]
270 ).reshape(2, 2)
271 tmp_pixels = tmp_array
273 # now for each 512x512 sub pixel region write the hips image with the corresponding healpix id
274 hpx_id_array = tmp_pixels
275 for i in range(binned_array.shape[0] // 512):
276 for j in range(binned_array.shape[1] // 512):
277 pixel_id = int(hpx_id_array[i, j])
278 sub_pixel = binned_array[i * 512 : i * 512 + 512, j * 512 : j * 512 + 512, :]
279 self.log.info(f"writing sub_pixel {pixel_id}")
280 _write_hips_image(
281 sub_pixel,
282 pixel_id,
283 hips_level,
284 self.hips_base_path,
285 self.config.file_extension,
286 self.config.array_type,
287 )
289 # Finally, bin the level 8 hpx to 256x256 (1/4 order 7) to save to the butler.
290 # This makes smaller arrays to load, and saves the binning operation in the joint phase.
291 zoomed = cv2.resize(output_array_hpx, (256, 256), interpolation=cv2.INTER_LANCZOS4)
293 return Struct(output_hpx=zoomed)
295 def _make_box_mask(self, shape: tuple[int, int], patch_grow: int) -> NDArray:
296 """Create a feathered box to use in blending images together.
298 This is used when mixing patches together into a larger contiguous region
299 such that overlaps are blended by a ramp function. This will be 50%
300 for each area at exactly halfway through the overlap, and 1/0 at the
301 extremes.
303 Parameters
304 ----------
305 shape : `tuple` of `int`, `int`
306 Shape of the output weighting mask.
307 patch_grow : `int`
308 Amount to grow bounding boxes.
310 Returns
311 -------
312 mask : `np.ndarray`
313 The output weighting mask to use in blending, repeated along third
314 axis to use with RGB images.
316 """
317 # Ramp in the overlap regions for patches as defined by the patch_grow factor
318 yind, xind = np.mgrid[: patch_grow * 4, : patch_grow * 4]
319 dis = ((yind - 2 * patch_grow) ** 2 + (xind - 2 * patch_grow) ** 2) ** 0.5
320 radial = 1 - np.clip(dis / (2 * patch_grow), 0, 1)
322 ramp = np.linspace(0, 1, patch_grow * 2)
323 tmpImg = np.zeros(shape)
324 tmpImg[: patch_grow * 2, :] = np.repeat(np.expand_dims(ramp, 1), tmpImg.shape[1], axis=1)
326 tmpImg[-1 * patch_grow * 2 :, :] = np.repeat( # noqa: E203
327 np.expand_dims(1 - ramp, 1), tmpImg.shape[1], axis=1
328 )
329 tmpImg[:, : patch_grow * 2] = np.repeat(np.expand_dims(ramp, 0), tmpImg.shape[0], axis=0)
331 tmpImg[:, -1 * patch_grow * 2 :] = np.repeat( # noqa: E203
332 np.expand_dims(1 - ramp, 0), tmpImg.shape[0], axis=0
333 )
334 # fix the corners
335 tmpImg[: 2 * patch_grow, : 2 * patch_grow] = radial[: 2 * patch_grow, : 2 * patch_grow]
336 tmpImg[: 2 * patch_grow, -2 * patch_grow :] = radial[: 2 * patch_grow, -2 * patch_grow :]
337 tmpImg[-2 * patch_grow :, : 2 * patch_grow] = radial[-2 * patch_grow :, : 2 * patch_grow]
338 tmpImg[-2 * patch_grow :, -2 * patch_grow :] = radial[-2 * patch_grow :, -2 * patch_grow :]
339 tmpImg = np.repeat(np.expand_dims(tmpImg, 2), 3, axis=2)
340 return tmpImg
342 def _assemble_sub_region(
343 self, tract_patch: dict[int, Iterable[tuple[DeferredDatasetHandle, SkyWcs, Box2I]]], patch_grow: int
344 ) -> list[tuple[NDArray, SkyWcs, Box2I]]:
345 """Assemble all the patches in each tract into images.
347 This function takes in an input keyed by tract, with values
348 corresponding the patches in that tract that overlap the quatum's
349 healpix value. It assembles each of these into a single image such
350 that the return values is a list of images (and metadata) one element
351 for each input tract.
353 Parameters
354 ----------
355 tract_patch : `dict` of `int` to `iterable` of `tuple` of
356 `DeferredDatasetHandle`, `SkyWcs` and `Box2I`
357 Input images and metadata organized into corresponding tracts.
358 patch_grow : `int`
359 Amount to grow patches by
361 Returns
362 -------
363 output_list : `list` of `tuple` of `NDArray` `SkyWcs` and `Box2I`
364 List of assembled images and metadata, one element for each tract
366 """
367 tmpImg = None
369 boxes = []
370 for _, iterable in tract_patch.items():
371 new_box = Box2I()
372 for _, _, bbox in iterable:
373 new_box.include(bbox)
374 # allocate tmp array
375 new_array = np.zeros((new_box.getHeight(), new_box.getWidth(), 3), dtype=np.float32)
376 for handle, skyWcs, box in iterable:
377 # Make a new box of the same size, but with the origin centered
378 # on the lowest corner were there is data.
379 localOrigin = box.getBegin() - new_box.getBegin()
380 localOrigin = Point2I(
381 x=int(np.floor(localOrigin.x)),
382 y=int(np.floor(localOrigin.y)),
383 )
385 localExtent = Extent2I(
386 x=int(np.floor(box.getWidth())),
387 y=int(np.floor(box.getHeight())),
388 )
389 tmpBox = Box2I(localOrigin, localExtent)
391 image = handle.get()
392 if tmpImg is None:
393 tmpImg = self._make_box_mask(image.shape[:2], patch_grow)
395 # Find all the pixels that are already populated
396 existing = ~np.all(new_array[*tmpBox.slices] == 0, axis=2)
398 # Populate all the pixels that don't already have a value in them
399 new_array[*tmpBox.slices][~existing, :] = image[~existing, :]
401 # For pixels with an existing value, compute the weighted average
402 # of values and assign that.
403 new_array[*tmpBox.slices][existing, :] = (
404 tmpImg[existing] * image[existing]
405 + (1 - tmpImg[existing]) * new_array[*tmpBox.slices][existing, :]
406 )
407 boxes.append((new_array, skyWcs, new_box))
408 return boxes
410 def runQuantum(
411 self,
412 butlerQC: QuantumContext,
413 inputRefs: InputQuantizedConnection,
414 outputRefs: OutputQuantizedConnection,
415 ) -> None:
416 # First get what healpix pixel this task is working on
417 healpix_id = butlerQC.quantum.dataId["healpix8"]
419 # grab the skymap
420 skymap: BaseSkyMap = butlerQC.get(inputRefs.skymap)
422 # Iterate over the input image refs, to get the corresponding bbox
423 # and assemble into container for run
424 inputs_by_tract = {}
425 for input_image_ref in inputRefs.input_images:
426 tract = input_image_ref.dataId["tract"]
427 patch = input_image_ref.dataId["patch"]
428 # All boxes in a given skymap will have the same inner dimensions
429 # for x and y and will be the same for all patches
430 imageWcs = skymap[tract][patch].getWcs()
431 box = skymap[tract][patch].getInnerBBox()
432 patch_grow = skymap[tract][patch].getCellInnerDimensions().getX()
433 box = box.dilatedBy(patch_grow)
434 imageHandle = butlerQC.get(input_image_ref)
435 container = inputs_by_tract.setdefault(tract, list())
436 container.append((imageHandle, imageWcs, box))
438 input_images = self._assemble_sub_region(inputs_by_tract, patch_grow)
440 outputs = self.run(input_images, healpix_id)
441 butlerQC.put(outputs, outputRefs)