Coverage for python/lsst/dax/images/cutout/_image_cutout.py: 64%
306 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-07-08 02:23 -0700
1# This file is part of dax_images_cutout.
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/>.
22from __future__ import annotations
24__all__ = ("CutoutMode", "Extraction", "ImageCutoutFactory")
26import dataclasses
27import logging
28from collections.abc import Sequence
29from enum import Enum, auto
30from typing import TYPE_CHECKING
31from uuid import UUID, uuid4
33import astropy.io.fits
34import astropy.time
36import lsst.images
37import lsst.images.serialization
38from lsst.daf.butler import Butler, DataId, DatasetRef
39from lsst.images import Box, MaskPlane, MaskSchema
40from lsst.images import Mask as ImagesMask
41from lsst.images.fits import DEFAULT_PAGE_SIZE, READ_CACHE_MAX_BYTES, ExtensionKey, FitsOpaqueMetadata
42from lsst.images.fits._input_archive import _READ_CACHE_TYPE
43from lsst.resources import ResourcePath, ResourcePathExpression
44from lsst.utils.timer import time_this
46from ._fits_projection import projection_and_bbox_from_fits_header
47from .projection_finders import ProjectionFinder
48from .stencils import PixelStencil, SkyStencil
49from .version import __version__
51if TYPE_CHECKING:
52 from lsst.afw.image import Exposure, Image, Mask, MaskedImage
53 from lsst.daf.base import PropertyList
56# Default logger.
57_LOG = logging.getLogger(__name__)
58_TIMER_LOG_LEVEL = logging.INFO
61class CutoutMode(Enum):
62 # The entire Exposure plus all associated metadata.
63 FULL_EXPOSURE = auto()
64 # Retrieve the full Exposure but only return MaskedImage/SkyWcs/metadata
65 STRIPPED_EXPOSURE = auto()
66 # Retrieve and return Image with no WCS or extra metadata.
67 IMAGE_ONLY = auto()
68 # MaskedImage using butler+afw
69 MASKED_IMAGE = auto()
70 # Astropy Image + primary HDU + WCS
71 ASTROPY_IMAGE = auto()
72 # Astropy with masked image.
73 ASTROPY_MASKED_IMAGE = auto()
76def _property_list_from_header(header: astropy.io.fits.Header) -> PropertyList:
77 """Convert an `astropy.io.fits.Header` to a `~lsst.daf.base.PropertyList`,
78 preserving each card's comment.
79 """
80 from lsst.daf.base import PropertyList
82 metadata = PropertyList()
83 for card in header.cards:
84 metadata.set(card.keyword, card.value, card.comment)
85 return metadata
88@dataclasses.dataclass
89class Extraction:
90 """A struct that aggregates the results of extracting subimages in the
91 image cutout backend.
92 """
94 cutout: Image | Mask | MaskedImage | Exposure | lsst.images.GeneralizedImage
95 """The image cutout itself.
96 """
98 sky_stencil: SkyStencil
99 """The original sky-coordinate stencil.
100 """
102 pixel_stencil: PixelStencil
103 """A pixel-coordinate representation of the stencil.
104 """
106 metadata: astropy.io.fits.Header
107 """Provenance FITS header cards describing the cutout process.
109 Written to the primary header on output. For afw cutout types it is
110 converted to a `~lsst.daf.base.PropertyList` first.
111 """
113 origin_ref: DatasetRef
114 """Fully-resolved reference to the dataset the cutout is from.
115 """
117 def mask(self, name: str = "OUTSIDE_STENCIL") -> None:
118 """Set the bitmask to show the approximate coverage of nonrectangular
119 stencils.
121 Parameters
122 ----------
123 name : `str`, optional
124 Name of the mask plane to add and set.
126 Notes
127 -----
128 This method modifies `cutout` in place if it is a `Mask`,
129 `MaskedImage`, or `Exposure`. It does nothing if `cutout` is an
130 `Image`.
131 """
132 if isinstance(self.cutout, lsst.images.MaskedImage):
133 mask = self.cutout.mask
134 # Create the new plane if it's not there.
135 if name not in mask.schema.names: 135 ↛ 139line 135 didn't jump to line 139 because the condition on line 135 was always true
136 # Adding a plane creates a whole new mask.
137 mask = mask.add_plane(name, "Pixel lies outside the stencil")
138 self.cutout.mask = mask
139 self.pixel_stencil.set_mask(mask, name, covered=False)
140 self._record_legacy_mask_plane(self.cutout, name)
141 return
143 # Try afw variants. Protect the imports (if the imports fail it is
144 # not possible for it to be an afw type).
145 try:
146 from lsst.afw.image import Exposure, Mask, MaskedImage
148 match self.cutout:
149 case Exposure() | MaskedImage(): 149 ↛ 150line 149 didn't jump to line 150 because the pattern on line 149 never matched
150 mask = self.cutout.mask
151 case Mask(): 151 ↛ 152line 151 didn't jump to line 152 because the pattern on line 151 never matched
152 mask = self.cutout
153 case _:
154 return
155 except ImportError:
156 # Can not determine the type so do not mask anything.
157 return
159 # Stage the coverage in an lsst.images.Mask, then OR it into the afw
160 # mask plane so stencils never sees an afw object.
161 images_mask = ImagesMask(
162 schema=MaskSchema([MaskPlane(name, "Pixel lies outside the stencil")]),
163 bbox=Box.from_legacy(mask.getBBox()),
164 )
165 self.pixel_stencil.set_mask(images_mask, name, covered=False)
166 outside = images_mask.get(name)
167 mask.addMaskPlane(name)
168 bits = mask.getPlaneBitMask(name)
169 mask.array[:, :] |= (bits * outside).astype(mask.array.dtype)
171 def _record_legacy_mask_plane(self, masked_image: lsst.images.MaskedImage, name: str) -> None:
172 """Add an ``MP_`` card for a newly added mask plane if the cutout came
173 from a legacy afw file.
175 Parameters
176 ----------
177 masked_image : `lsst.images.MaskedImage`
178 The cutout whose mask just gained the ``name`` plane.
179 name : `str`
180 Name of the mask plane just added to ``masked_image.mask``.
182 Notes
183 -----
184 ``MaskedImage.from_hdu_list`` keeps the legacy ``MP_*`` mask-plane
185 cards (re-indexed to the reshuffled schema) in the ``MASK`` opaque
186 metadata so ``lsst.afw.image`` can still read the inherited planes. A
187 plane added afterwards (e.g. ``OUTSIDE_STENCIL``) is only described by
188 the native ``MSKN`` cards, which afw does not understand, so afw would
189 not see it. When the cutout carries that legacy provenance, record a
190 matching ``MP_`` card for the new plane at the bit it occupies in the
191 serialized (single-element) schema, so the whole cutout mask stays
192 readable by legacy tooling. Files without legacy ``MP_`` cards are
193 left untouched.
195 Reaching into ``_opaque_metadata`` mirrors how the cutout's ``TIMESYS``
196 is recovered elsewhere in this module.
197 """
198 opaque = masked_image._opaque_metadata
199 if not isinstance(opaque, FitsOpaqueMetadata):
200 return
201 key = ExtensionKey("MASK")
202 mask_header = opaque.headers.get(key)
203 if mask_header is None or not any(keyword.startswith("MP_") for keyword in mask_header): 203 ↛ 207line 203 didn't jump to line 207 because the condition on line 203 was always true
204 return
205 # The serialized form packs every plane into one element, so a plane's
206 # bit index is its position among the non-placeholder planes.
207 serialized_bit = 0
208 for plane in masked_image.mask.schema:
209 if plane is None:
210 continue
211 if plane.name == name:
212 break
213 serialized_bit += 1
214 else:
215 return
216 # Opaque headers are treated as immutable, so replace, don't mutate.
217 updated = mask_header.copy()
218 updated[f"MP_{name}"] = serialized_bit
219 opaque.headers[key] = updated
221 def write_fits(self, path: str, logger: logging.Logger | None = None) -> None:
222 """Write the cutout to a FITS file.
224 Parameters
225 ----------
226 path : `str`
227 Local path to the file.
228 logger : `logging.Logger`, optional
229 Logger to use for timing messages. If `None`, a default logger
230 will be used.
232 Notes
233 -----
234 If ``cutout`` is an `Exposure`, this will merge `metadata` into the
235 cutout's attached metadata. In other cases, `metadata` is written
236 to the primary header without modifying `cutout`.
237 """
238 logger = logger if logger is not None else _LOG
239 with time_this(logger, msg="Writing FITS file to %s", args=(path,), level=_TIMER_LOG_LEVEL):
240 if isinstance(self.cutout, lsst.images.GeneralizedImage): 240 ↛ 247line 240 didn't jump to line 247 because the condition on line 240 was always true
241 # Write the cutout provenance into the primary FITS header via
242 # the archive's update_header hook. ``metadata`` is an astropy
243 # Header, so card comments are preserved. Storing it in the
244 # object's flexible metadata would instead place it in the JSON
245 # tree, where it would not appear in the primary header.
246 self.cutout.write(path, update_header=lambda header: header.update(self.metadata))
247 elif hasattr(self.cutout, "getMetadata") and hasattr(self.cutout, "writeFits"):
248 self.cutout.metadata.update(_property_list_from_header(self.metadata))
249 self.cutout.writeFits(path)
250 else:
251 self.cutout.writeFits(path, metadata=_property_list_from_header(self.metadata))
254class ImageCutoutFactory:
255 """High-level interface to the image cutout functionality.
257 Parameters
258 ----------
259 butler : `Butler`
260 Butler that subimages are extracted from.
261 projection_finder : `ProjectionObject`
262 Object that obtains the WCS and bounding box for butler datasets of
263 different types. May include caches.
264 output_root : convertible to `ResourcePath`
265 Root of output file URIs.
266 temporary_root : convertible to `ResourcePath`, optional
267 Local filesystem root to write files to before they are transferred to
268 ``output_root`` (passed as the prefix argument to
269 `ResourcePath.temporary_uri`).
270 logger : `logging.Logger`, optional
271 Logger to use for timing messages. If `None`, a default logger
272 will be used.
273 """
275 def __init__(
276 self,
277 butler: Butler,
278 projection_finder: ProjectionFinder,
279 output_root: ResourcePathExpression,
280 temporary_root: ResourcePathExpression | None = None,
281 logger: logging.Logger | None = None,
282 ):
283 self.butler = butler
284 self.projection_finder = projection_finder
285 self.output_root = ResourcePath(output_root, forceAbsolute=True, forceDirectory=True)
286 self.temporary_root = (
287 ResourcePath(temporary_root, forceDirectory=True) if temporary_root is not None else None
288 )
289 self.logger = logger if logger is not None else _LOG
291 butler: Butler
292 """Butler that subimage are extracted from (`Butler`).
293 """
295 projection_finder: ProjectionFinder
296 """Object that obtains the WCS and bounding box for butler datasets of
297 different types (`ProjectionFinder`).
298 """
300 output_root: ResourcePath
301 """Root path that extracted cutouts are written to (`ResourcePath`).
302 """
304 temporary_root: ResourcePath | None
305 """Local filesystem root to write files to before they are transferred to
306 ``output_root``
307 """
309 def process_ref(
310 self,
311 stencil: SkyStencil,
312 ref: DatasetRef,
313 *,
314 mask_plane: str | None = "OUTSIDE_STENCIL",
315 cutout_mode: CutoutMode = CutoutMode.FULL_EXPOSURE,
316 ) -> ResourcePath:
317 """Extract and write a cutout from a fully-resolved `DatasetRef`.
319 Parameters
320 ----------
321 stencil : `SkyStencil`
322 Definition of the cutout region, in sky coordinates.
323 ref : `DatasetRef`
324 Fully-resolved reference to the dataset to obtain the cutout from.
325 Must have ``DatasetRef.id`` not `None` (use `extract_search`
326 instead when this is not the case). Need not have an expanded data
327 ID. May represent an image-like dataset component.
328 mask_plane : `str`, optional
329 If not `None`, set this mask plane in the extracted cutout to flag
330 pixels that lie outside the stencil region. Does nothing if the
331 image type does not have a mask plane. Defaults to
332 ``OUTSIDE_STENCIL``.
334 Returns
335 -------
336 uri : `ResourcePath`
337 Full path to the extracted cutout.
338 """
339 extract_result = self.extract_ref(stencil, ref, cutout_mode=cutout_mode)
340 if mask_plane is not None: 340 ↛ 342line 340 didn't jump to line 342 because the condition on line 340 was always true
341 extract_result.mask(mask_plane)
342 return self.write_fits(extract_result)
344 def process_uuid(
345 self,
346 stencil: SkyStencil,
347 uuid: UUID,
348 *,
349 component: str | None = None,
350 mask_plane: str | None = "OUTSIDE_STENCIL",
351 cutout_mode: CutoutMode = CutoutMode.FULL_EXPOSURE,
352 ) -> ResourcePath:
353 """Extract and write a cutout from a dataset identified by its UUID.
355 Parameters
356 ----------
357 stencil : `SkyStencil`
358 Definition of the cutout region, in sky coordinates.
359 uuid : `UUID`
360 Unique ID of the dataset to extract the subimage from.
361 component : `str`, optional
362 If not `None` (default), read this component instead of the
363 composite dataset.
364 mask_plane : `str`, optional
365 If not `None`, set this mask plane in the extracted cutout to flag
366 pixels that lie outside the stencil region. Does nothing if the
367 image type does not have a mask plane. Defaults to
368 ``OUTSIDE_STENCIL``.
370 Returns
371 -------
372 uri : `ResourcePath`
373 Full path to the extracted cutout.
374 """
375 extract_result = self.extract_uuid(stencil, uuid, component=component, cutout_mode=cutout_mode)
376 if mask_plane is not None: 376 ↛ 378line 376 didn't jump to line 378 because the condition on line 376 was always true
377 extract_result.mask(mask_plane)
378 return self.write_fits(extract_result)
380 def process_search(
381 self,
382 stencil: SkyStencil,
383 dataset_type_name: str,
384 data_id: DataId,
385 collections: Sequence[str],
386 *,
387 mask_plane: str | None = "OUTSIDE_STENCIL",
388 ) -> ResourcePath:
389 """Extract and write a cutout from a dataset identified by a
390 (dataset type, data ID, collection path) tuple.
392 Parameters
393 ----------
394 stencil : `SkyStencil`
395 Definition of the cutout region, in sky coordinates.
396 dataset_type_name : `str`
397 Name of the butler dataset. Use ``.``-separate terms to read an
398 image-like component.
399 data_id : `dict` or `DataCoordinate`
400 Mapping-of-dimensions identifier for the dataset within its
401 collection.
402 collections : `Iterable` [ `str` ]
403 Collections to search for the dataset, in the order they should be
404 searched.
405 mask_plane : `str`, optional
406 If not `None`, set this mask plane in the extracted cutout to flag
407 pixels that lie outside the stencil region. Does nothing if the
408 image type does not have a mask plane. Defaults to
409 ``OUTSIDE_STENCIL``.
411 Returns
412 -------
413 uri : `ResourcePath`
414 Full path to the extracted cutout.
415 """
416 extract_result = self.extract_search(stencil, dataset_type_name, data_id, collections)
417 if mask_plane is not None:
418 extract_result.mask(mask_plane)
419 return self.write_fits(extract_result)
421 def extract_ref(
422 self, stencil: SkyStencil, ref: DatasetRef, cutout_mode: CutoutMode = CutoutMode.FULL_EXPOSURE
423 ) -> Extraction:
424 """Extract a subimage from a fully-resolved `DatasetRef`.
426 Parameters
427 ----------
428 stencil : `SkyStencil`
429 Definition of the cutout region, in sky coordinates.
430 ref : `DatasetRef`
431 Fully-resolved reference to the dataset to obtain the cutout from.
432 Must have ``DatasetRef.id`` not `None` (use `extract_search`
433 instead when this is not the case). Need not have an expanded data
434 ID. May represent an image-like dataset component.
436 Returns
437 -------
438 result : `Extraction`
439 Struct that combines the cutout itself with additional metadata
440 and the pixel-coordinate stencil. The cutout is not masked;
441 `Extraction.mask` must be called explicitly if desired.
442 """
443 if issubclass(ref.datasetType.storageClass.pytype, lsst.images.GeneralizedImage): 443 ↛ 445line 443 didn't jump to line 445 because the condition on line 443 was always true
444 return self._extract_ref_v2(stencil, ref, cutout_mode=cutout_mode)
445 return self._extract_ref_legacy(stencil, ref, cutout_mode=cutout_mode)
447 def _read_astropy_hdulist(
448 self,
449 cutout_mode: CutoutMode,
450 stencil: SkyStencil,
451 ref: DatasetRef,
452 ) -> tuple[lsst.images.GeneralizedImage, PixelStencil | None, str]:
453 """Read the primary header and pixel HDUs, cut to the stencil.
455 Parameters
456 ----------
457 cutout_mode
458 The requested cutout mode. Must be either an Astropy image or
459 masked image request.
460 stencil
461 Sky-coordinate stencil defining the cutout region.
462 ref : `DatasetRef`
463 Resolved reference to the dataset to read.
465 Returns
466 -------
467 result : `lsst.images.GeneralizedImage`
468 The resultant generalized image constructed from the individual
469 HDUs.
470 pixel_stencil : `PixelStencil` or `None`
471 Pixel-coordinate stencil computed from the first pixel HDU, or
472 `None` if no pixel HDU was found.
473 timesys : `str`
474 ``TIMESYS`` from the primary header, or ``"UTC"``.
475 """
476 if cutout_mode not in (CutoutMode.ASTROPY_IMAGE, CutoutMode.ASTROPY_MASKED_IMAGE): 476 ↛ 477line 476 didn't jump to line 477 because the condition on line 476 was never true
477 raise ValueError(f"Unexpected cutout mode {cutout_mode} encountered")
479 pixel_components = (
480 {"image"} if cutout_mode == CutoutMode.ASTROPY_IMAGE else {"image", "mask", "variance"}
481 )
482 # Tune the fsspec cache to match what we use in lsst.images.
483 maxblocks = max(2, READ_CACHE_MAX_BYTES // DEFAULT_PAGE_SIZE)
484 fsspec_kwargs = {
485 "block_size": DEFAULT_PAGE_SIZE,
486 "cache_type": _READ_CACHE_TYPE,
487 "cache_options": {"maxblocks": maxblocks},
488 }
489 timesys = "UTC"
490 pixel_stencil: PixelStencil | None = None
491 bbox: Box | None = None
492 uri = self.butler.getURI(ref)
493 fs, fspath = uri.to_fsspec()
494 hdul: list[astropy.io.fits.hdu.base.ExtensionHDU] = []
495 found_primary = False
496 with fs.open(fspath, **fsspec_kwargs) as f, astropy.io.fits.open(f) as fits_obj:
497 for hdu in fits_obj: 497 ↛ 553line 497 didn't jump to line 553
498 if not found_primary:
499 hdul.append(hdu.copy())
500 timesys_hdr = hdul[0].header.get("TIMESYS", timesys)
501 if timesys_hdr: 501 ↛ 506line 501 didn't jump to line 506 because the condition on line 501 was always true
502 # For mypy since in theory a FITS header can exist
503 # but be undefined.
504 timesys = str(timesys_hdr)
506 found_primary = True
508 # Some old legacy data erroneously wrote "IMAGE" to the
509 # EXTNAME of the primary header. Since we are assuming
510 # here that the primary never has pixel data, fix the
511 # header to prevent downstream confusion.
512 if "EXTNAME" in hdul[0].header: 512 ↛ 513line 512 didn't jump to line 513 because the condition on line 512 was never true
513 del hdul[0].header["EXTNAME"]
515 continue
517 hdr = hdu.header
518 extname = hdr.get("EXTNAME")
519 if extname and extname.lower() in pixel_components: 519 ↛ 550line 519 didn't jump to line 550 because the condition on line 519 was always true
520 pixel_components.remove(extname.lower())
521 projection, full_bbox = projection_and_bbox_from_fits_header(hdr, hdu.shape)
522 if pixel_stencil is None:
523 pixel_stencil = stencil.to_pixels(projection, full_bbox)
524 bbox = pixel_stencil.bbox
526 assert bbox is not None
527 # Offsets of the cutout within the full HDU (array order).
528 min_x = bbox.x.start - full_bbox.x.start
529 max_x = bbox.x.stop - full_bbox.x.start
530 min_y = bbox.y.start - full_bbox.y.start
531 max_y = bbox.y.stop - full_bbox.y.start
532 data = hdu.section[min_y:max_y, min_x:max_x].copy()
534 # Correct the header WCS for the cutout offset.
535 if (k := "CRPIX1") in hdr: 535 ↛ 537line 535 didn't jump to line 537 because the condition on line 535 was always true
536 hdr[k] -= min_x
537 if (k := "CRPIX2") in hdr: 537 ↛ 539line 537 didn't jump to line 539 because the condition on line 537 was always true
538 hdr[k] -= min_y
539 if (k := "LTV1") in hdr: 539 ↛ 540line 539 didn't jump to line 540 because the condition on line 539 was never true
540 hdr[k] = -bbox.x.start
541 if (k := "LTV2") in hdr: 541 ↛ 542line 541 didn't jump to line 542 because the condition on line 541 was never true
542 hdr[k] = -bbox.y.start
543 if (k := "CRVAL1A") in hdr: 543 ↛ 545line 543 didn't jump to line 545 because the condition on line 543 was always true
544 hdr[k] = bbox.x.start
545 if (k := "CRVAL2A") in hdr: 545 ↛ 548line 545 didn't jump to line 548 because the condition on line 545 was always true
546 hdr[k] = bbox.y.start
548 hdul.append(astropy.io.fits.ImageHDU(data=data, header=hdr.copy()))
550 if not pixel_components:
551 break
553 hdulist = astropy.io.fits.HDUList(hdus=hdul)
554 result: lsst.images.GeneralizedImage
555 if cutout_mode == CutoutMode.ASTROPY_IMAGE:
556 result = lsst.images.Image.from_hdu_list(hdulist)
557 else:
558 result = lsst.images.MaskedImage.from_hdu_list(hdulist)
560 return result, pixel_stencil, timesys
562 def _extract_ref_legacy(
563 self, stencil: SkyStencil, ref: DatasetRef, cutout_mode: CutoutMode = CutoutMode.FULL_EXPOSURE
564 ) -> Extraction:
565 """Extract a subimage from a fully-resolved `DatasetRef` associated
566 with a legacy exposure.
568 Parameters
569 ----------
570 stencil : `SkyStencil`
571 Definition of the cutout region, in sky coordinates.
572 ref : `DatasetRef`
573 Fully-resolved reference to the dataset to obtain the cutout from.
574 Must have ``DatasetRef.id`` not `None` (use `extract_search`
575 instead when this is not the case). Need not have an expanded data
576 ID. May represent an image-like dataset component.
578 Returns
579 -------
580 result : `Extraction`
581 Struct that combines the cutout itself with additional metadata
582 and the pixel-coordinate stencil. The cutout is not masked;
583 `Extraction.mask` must be called explicitly if desired.
584 """
585 # We know that afw is available here since we received an afw Exposure.
586 from lsst.afw.image import makeExposure, makeMaskedImage
588 if ref.id is None:
589 raise ValueError(f"A resolved DatasetRef is required; got {ref}.")
590 # Timestamp of the cutout extraction.
591 now = astropy.time.Time.now()
592 timesys = "UTC"
593 # Get the WCS and bbox of this dataset unless we are in astropy mode.
594 wcs = None
595 bbox = None
596 pixel_stencil = None
597 if cutout_mode not in (CutoutMode.ASTROPY_IMAGE, CutoutMode.ASTROPY_MASKED_IMAGE):
598 projection, bbox = self.projection_finder(ref, self.butler, logger=self.logger)
599 # Transform the stencil to pixel coordinates.
600 pixel_stencil = stencil.to_pixels(projection, bbox)
602 # Actually read the cutout. Leave it to the butler to cache remote
603 # files locally or do partial remote reads.
604 with time_this(
605 self.logger,
606 msg="Extract cutout",
607 kwargs={"id": str(ref.id), "cutout_mode": str(cutout_mode), "stencil": str(stencil)},
608 level=_TIMER_LOG_LEVEL,
609 ):
610 match cutout_mode:
611 case CutoutMode.FULL_EXPOSURE:
612 assert pixel_stencil is not None
613 cutout = self.butler.get(ref, parameters={"bbox": pixel_stencil.bbox.to_legacy()})
614 timesys = cutout.metadata.get("TIMESYS", timesys)
615 case CutoutMode.STRIPPED_EXPOSURE:
616 assert pixel_stencil is not None
617 cutout = self.butler.get(ref, parameters={"bbox": pixel_stencil.bbox.to_legacy()})
618 original_metadata = cutout.metadata
619 timesys = original_metadata.get("TIMESYS", timesys)
620 cutout = makeExposure(cutout.maskedImage, wcs=cutout.wcs)
622 # A full exposure so we can store the metadata back in it.
623 cutout.setMetadata(original_metadata)
624 case CutoutMode.IMAGE_ONLY:
625 assert pixel_stencil is not None
626 cutout = self.butler.get(
627 ref.makeComponentRef("image"), parameters={"bbox": pixel_stencil.bbox.to_legacy()}
628 )
629 # No metadata so UTC is default.
630 timesys = "UTC"
631 case CutoutMode.MASKED_IMAGE:
632 # Rely on the file being cached on first read. Faster than
633 # reading entire exposure.
634 assert pixel_stencil is not None
635 image = self.butler.get(
636 ref.makeComponentRef("image"), parameters={"bbox": pixel_stencil.bbox.to_legacy()}
637 )
638 variance = self.butler.get(
639 ref.makeComponentRef("variance"), parameters={"bbox": pixel_stencil.bbox.to_legacy()}
640 )
641 mask = self.butler.get(
642 ref.makeComponentRef("mask"), parameters={"bbox": pixel_stencil.bbox.to_legacy()}
643 )
644 wcs = self.butler.get(ref.makeComponentRef("wcs"))
645 masked_image = makeMaskedImage(image, mask, variance)
646 cutout = makeExposure(masked_image, wcs=wcs)
648 original_metadata = self.butler.get(ref.makeComponentRef("metadata"))
649 timesys = original_metadata.get("TIMESYS", timesys)
651 # The cutout is an Exposure so we can attach metadata.
652 cutout.setMetadata(original_metadata)
654 case CutoutMode.ASTROPY_IMAGE | CutoutMode.ASTROPY_MASKED_IMAGE:
655 # Bypass butler and read the pixel HDU directly.
656 cutout, pixel_stencil, timesys = self._read_astropy_hdulist(cutout_mode, stencil, ref)
657 case _:
658 raise ValueError(f"Unsupported cutout mode: {cutout_mode}")
660 # Create the provenance header with the cutout parameters.
661 metadata = self._record_cutout_provenance(ref, now, stencil, timesys)
663 # Every supported cutout mode produces a pixel stencil above.
664 assert pixel_stencil is not None
665 return Extraction(
666 cutout=cutout,
667 sky_stencil=stencil,
668 pixel_stencil=pixel_stencil,
669 metadata=metadata,
670 origin_ref=ref,
671 )
673 def _extract_ref_v2(
674 self, stencil: SkyStencil, ref: DatasetRef, cutout_mode: CutoutMode = CutoutMode.FULL_EXPOSURE
675 ) -> Extraction:
676 """Extract a subimage from a fully-resolved `DatasetRef` associated
677 with a lsst.images model.
679 Parameters
680 ----------
681 stencil : `SkyStencil`
682 Definition of the cutout region, in sky coordinates.
683 ref : `DatasetRef`
684 Fully-resolved reference to the dataset to obtain the cutout from.
685 Must have ``DatasetRef.id`` not `None` (use `extract_search`
686 instead when this is not the case). Need not have an expanded data
687 ID. May represent an image-like dataset component.
689 Returns
690 -------
691 result : `Extraction`
692 Struct that combines the cutout itself with additional metadata
693 and the pixel-coordinate stencil. The cutout is not masked;
694 `Extraction.mask` must be called explicitly if desired.
695 """
696 if ref.id is None: 696 ↛ 697line 696 didn't jump to line 697 because the condition on line 696 was never true
697 raise ValueError(f"A resolved DatasetRef is required; got {ref}.")
698 # Timestamp of the cutout extraction.
699 now = astropy.time.Time.now()
700 timesys = "UTC"
701 pixel_stencil = None
703 # There are two distinct modes of operation. One is using the native
704 # lsst.images interface, the other is using Astropy to deal with
705 # IMAGE and MASKED_IMAGE by using standard FITS conventions.
706 if cutout_mode not in (CutoutMode.ASTROPY_IMAGE, CutoutMode.ASTROPY_MASKED_IMAGE):
707 # To reduce round-trips to an object store we want to open the file
708 # once and then read out the components we need. In theory we can
709 # ask Butler get to return multiple components but if we do that
710 # we do not know the bounding box to use unless we add a skybox
711 # parameter to get that lets us pass in a stencil directly.
712 # This means we have to use the URI for direct access, bypassing
713 # butler get for now.
714 uri = self.butler.getURI(ref)
716 # Time the full open and retrieval.
717 with time_this(
718 self.logger,
719 msg="Extract cutout",
720 kwargs={"id": str(ref.id), "cutout_mode": str(cutout_mode), "stencil": str(stencil)},
721 level=_TIMER_LOG_LEVEL,
722 ):
723 with lsst.images.serialization.open(uri) as reader:
724 sky_projection = reader.get_component("sky_projection")
725 bbox = reader.get_component("bbox")
727 # Transform the stencil to pixel coordinates.
728 pixel_stencil = stencil.to_pixels(sky_projection, bbox)
729 modern_bbox = pixel_stencil.bbox
731 match cutout_mode:
732 case CutoutMode.FULL_EXPOSURE:
733 cutout = reader.read(bbox=modern_bbox)
734 case CutoutMode.IMAGE_ONLY:
735 cutout = reader.get_component("image", bbox=modern_bbox)
736 case CutoutMode.MASKED_IMAGE | CutoutMode.STRIPPED_EXPOSURE: 736 ↛ 741line 736 didn't jump to line 741 because the pattern on line 736 always matched
737 # A Stripped exposure is meaningless with the
738 # new models since MaskedImage now carries a
739 # WCS and metadata.
740 cutout = reader.get_component("masked_image", bbox=modern_bbox)
741 case _:
742 raise ValueError(f"Unsupported cutout mode: {cutout_mode}")
744 if cutout._opaque_metadata is not None:
745 timesys = cutout._opaque_metadata.headers[ExtensionKey()].get("TIMESYS", timesys)
747 else:
748 # This is the Astropy direct branch.
749 with time_this(
750 self.logger,
751 msg="Extract cutout",
752 kwargs={"id": str(ref.id), "cutout_mode": str(cutout_mode), "stencil": str(stencil)},
753 level=_TIMER_LOG_LEVEL,
754 ):
755 # Bypass butler and lsst.images and go straight to the HDUs.
756 cutout, pixel_stencil, timesys = self._read_astropy_hdulist(cutout_mode, stencil, ref)
758 # Create some FITS metadata with the cutout parameters.
759 metadata = self._record_cutout_provenance(ref, now, stencil, timesys)
761 # Every supported cutout mode produces a pixel stencil above.
762 assert pixel_stencil is not None
763 return Extraction(
764 cutout=cutout,
765 sky_stencil=stencil,
766 pixel_stencil=pixel_stencil,
767 metadata=metadata,
768 origin_ref=ref,
769 )
771 def _record_cutout_provenance(
772 self, ref: DatasetRef, start_time: astropy.time.Time, stencil: SkyStencil, timesys: str
773 ) -> astropy.io.fits.Header:
774 header = astropy.io.fits.Header()
776 # Create some FITS metadata with the cutout parameters.
777 # Some of these are added as provenance by the butler on write so may
778 # no longer be necessary.
779 header.set("BTLRUUID", ref.id.hex, "Butler UUID of full image")
780 header.set("BTLRNAME", ref.datasetType.name, "Butler dataset type")
781 for n, (k, v) in enumerate(ref.dataId.required.items()):
782 # Write data ID dictionary sort of like a list of 2-tuples, to make
783 # it easier to stay within the FITS 8-char key limit.
784 header.set(f"BTLRK{n:03}", k, f"Name of dimension {n} in the data ID")
785 header.set(f"BTLRV{n:03}", v, f"Value of dimension {n} in the data ID")
786 header.extend(stencil.to_fits_metadata())
788 # Record the time and software version.
789 start_time.format = "fits"
790 start_time = start_time.tai if timesys.lower() == "tai" else start_time.utc
791 header.set("DATE-CUT", str(start_time), "Time of cutout extraction")
792 header.set("CUTVERS", __version__, "dax_images_cutout software version")
794 return header
796 def extract_uuid(
797 self,
798 stencil: SkyStencil,
799 uuid: UUID,
800 *,
801 component: str | None = None,
802 cutout_mode: CutoutMode = CutoutMode.FULL_EXPOSURE,
803 ) -> Extraction:
804 """Extract a subimage from a dataset identified by its UUID.
806 Parameters
807 ----------
808 stencil : `SkyStencil`
809 Definition of the cutout region, in sky coordinates.
810 uuid : `UUID`
811 Unique ID of the dataset to read from.
812 component : `str`, optional
813 If not `None` (default), read this component instead of the
814 composite dataset.
816 Returns
817 -------
818 result : `Extraction`
819 Struct that combines the cutout itself with additional metadata
820 and the pixel-coordinate stencil. The cutout is not masked;
821 `Extraction.mask` must be called explicitly if desired.
822 """
823 ref = self.butler.get_dataset(uuid)
824 if ref is None: 824 ↛ 825line 824 didn't jump to line 825 because the condition on line 824 was never true
825 raise LookupError(f"No dataset found with UUID {uuid}.")
826 if component is not None: 826 ↛ 827line 826 didn't jump to line 827 because the condition on line 826 was never true
827 ref = ref.makeComponentRef(component)
828 return self.extract_ref(stencil, ref, cutout_mode=cutout_mode)
830 def extract_search(
831 self, stencil: SkyStencil, dataset_type_name: str, data_id: DataId, collections: Sequence[str]
832 ) -> Extraction:
833 """Extract a subimage from a dataset identified by a (dataset type,
834 data ID, collection path) tuple.
836 Parameters
837 ----------
838 stencil : `SkyStencil`
839 Definition of the cutout region, in sky coordinates.
840 dataset_type_name : `str`
841 Name of the butler dataset. Use ``.``-separate terms to read an
842 image-like component.
843 data_id : `dict` or `DataCoordinate`
844 Mapping-of-dimensions identifier for the dataset within its
845 collection.
846 collections : `Iterable` [ `str` ]
847 Collections to search for the dataset, in the order they should be
848 searched.
850 Returns
851 -------
852 result : `Extraction`
853 Struct that combines the cutout itself with additional metadata
854 and the pixel-coordinate stencil. The cutout is not masked;
855 `Extraction.mask` must be called explicitly if desired.
856 """
857 ref = self.butler.find_dataset(dataset_type_name, data_id, collections=collections)
858 if ref is None:
859 raise LookupError(
860 f"No {dataset_type_name} dataset found with data ID {data_id} in {collections}."
861 )
862 return self.extract_ref(stencil, ref)
864 def write_fits(self, extract_result: Extraction) -> ResourcePath:
865 """Write a `Extraction` to a remote FITS file in `output_root`.
867 Parameters
868 ----------
869 extract_result : `Extraction`
870 Result of a call to a ``extract_*`` method.
872 Returns
873 -------
874 uri : `ResourcePath`
875 Full path to the extracted cutout.
876 """
877 output_uuid = uuid4()
878 remote_uri = self.output_root.join(output_uuid.hex + ".fits")
879 with ResourcePath.temporary_uri(prefix=self.temporary_root, suffix=".fits") as tmp_uri:
880 tmp_uri.parent().mkdir()
881 extract_result.write_fits(tmp_uri.ospath, logger=self.logger)
882 remote_uri.transfer_from(tmp_uri, transfer="copy")
883 return remote_uri