Coverage for python/lsst/images/_masked_image.py: 70%
192 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-07-16 15:24 -0700
1# This file is part of lsst-images.
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# Use of this source code is governed by a 3-clause BSD-style
10# license that can be found in the LICENSE file.
12from __future__ import annotations
14__all__ = ("MaskedImage", "MaskedImageSerializationModel")
16import functools
17from collections.abc import Mapping, MutableMapping
18from contextlib import ExitStack
19from types import EllipsisType
20from typing import TYPE_CHECKING, Any, ClassVar, Literal
22import astropy.io.fits
23import astropy.units
24import astropy.wcs
25import numpy as np
26import pydantic
28from lsst.resources import ResourcePath, ResourcePathExpression
30from . import fits
31from ._generalized_image import GeneralizedImage
32from ._geom import Box
33from ._image import DEFAULT_PIXEL_FRAME, Image, ImageSerializationModel
34from ._mask import Mask, MaskPlane, MaskSchema, MaskSerializationModel
35from ._transforms import Frame, SkyProjection, SkyProjectionSerializationModel
36from .serialization import (
37 ArchiveTree,
38 InputArchive,
39 InvalidParameterError,
40 MetadataValue,
41 OutputArchive,
42)
43from .utils import is_none
45if TYPE_CHECKING:
46 try:
47 from lsst.afw.image import MaskedImage as LegacyMaskedImage
48 except ImportError:
49 type LegacyMaskedImage = Any # type: ignore[no-redef]
52class MaskedImage(GeneralizedImage):
53 """A multi-plane image with data (image), mask, and variance planes.
55 Parameters
56 ----------
57 image
58 The main image plane. If this has a `SkyProjection`, it will be used
59 for all planes unless a ``sky_projection`` is passed separately.
60 mask
61 A bitmask image that annotates the main image plane. Must have the
62 same bounding box as ``image`` if provided. Any attached
63 ``sky_projection`` is replaced (possibly by `None`).
64 variance
65 The per-pixel uncertainty of the main image as an image of variance
66 values. Must have the same bounding box as ``image`` if provided, and
67 its units must be the square of ``image.unit`` or `None`.
68 Values default to ``1.0``. Any attached ``sky_projection`` is replaced
69 (possibly by `None`).
70 mask_schema
71 Schema for the mask plane. Must be provided if and only if ``mask`` is
72 not provided.
73 sky_projection
74 Projection that maps the pixel grid to the sky.
75 metadata
76 Arbitrary flexible metadata to associate with the image.
77 """
79 def __init__(
80 self,
81 image: Image,
82 *,
83 mask: Mask | None = None,
84 variance: Image | None = None,
85 mask_schema: MaskSchema | None = None,
86 sky_projection: SkyProjection[Any] | None = None,
87 metadata: dict[str, MetadataValue] | None = None,
88 ) -> None:
89 super().__init__(metadata)
90 if sky_projection is None:
91 sky_projection = image.sky_projection
92 else:
93 image = image.view(sky_projection=sky_projection)
94 if mask is None:
95 if mask_schema is None:
96 raise TypeError("'mask_schema' must be provided if 'mask' is not.")
97 mask = Mask(schema=mask_schema, bbox=image.bbox, sky_projection=sky_projection)
98 elif mask_schema is not None:
99 raise TypeError("'mask_schema' may not be provided if 'mask' is.")
100 else:
101 if image.bbox != mask.bbox:
102 raise ValueError(f"Image ({image.bbox}) and mask ({mask.bbox}) bboxes do not agree.")
103 mask = mask.view(sky_projection=sky_projection)
104 if variance is None:
105 variance = Image(
106 1.0,
107 dtype=np.float32,
108 bbox=image.bbox,
109 unit=None if image.unit is None else image.unit**2,
110 sky_projection=sky_projection,
111 )
112 else:
113 if image.bbox != variance.bbox:
114 raise ValueError(f"Image ({image.bbox}) and variance ({variance.bbox}) bboxes do not agree.")
115 variance = variance.view(sky_projection=sky_projection)
116 if image.unit is None:
117 if variance.unit is not None:
118 raise ValueError(f"Image has no units but variance does ({variance.unit}).")
119 elif variance.unit is None: 119 ↛ 120line 119 didn't jump to line 120 because the condition on line 119 was never true
120 variance = variance.view(unit=image.unit**2)
121 elif variance.unit != image.unit**2:
122 raise ValueError(
123 f"Variance unit ({variance.unit}) should be the square of the image unit ({image.unit})."
124 )
125 self._image = image
126 self._mask = mask
127 self._variance = variance
129 @property
130 def image(self) -> Image:
131 """The main image plane (`~lsst.images.Image`)."""
132 return self._image
134 @property
135 def mask(self) -> Mask:
136 """The mask plane (`~lsst.images.Mask`).
138 Assigning a new `~lsst.images.Mask` (for example one returned by
139 `~lsst.images.Mask.add_planes`) replaces the mask plane. The new mask
140 must share this image's bounding box; its sky projection is replaced
141 with the image's.
142 """
143 return self._mask
145 @mask.setter
146 def mask(self, value: Mask) -> None:
147 if value.bbox != self._image.bbox:
148 raise ValueError(f"Image ({self._image.bbox}) and mask ({value.bbox}) bboxes do not agree.")
149 if self._image.sky_projection != value.sky_projection: 149 ↛ 150line 149 didn't jump to line 150 because the condition on line 149 was never true
150 raise ValueError("Image sky projection and new mask sky projection do not agree")
151 # Use a view to ensure that the WCS instances across the masked image
152 # are the same projections.
153 self._mask = value.view(sky_projection=self._image.sky_projection)
155 @property
156 def variance(self) -> Image:
157 """The variance plane (`~lsst.images.Image`)."""
158 return self._variance
160 @property
161 def bbox(self) -> Box:
162 """The bounding box shared by all three image planes
163 (`~lsst.images.Box`).
164 """
165 return self._image.bbox
167 @property
168 def unit(self) -> astropy.units.UnitBase | None:
169 """The units of the image plane (`astropy.units.Unit` | `None`)."""
170 return self._image.unit
172 @property
173 def sky_projection(self) -> SkyProjection[Any] | None:
174 """The projection that maps the pixel grid to the sky
175 (`~lsst.images.SkyProjection` | `None`).
176 """
177 return self._image.sky_projection
179 def __getitem__(self, bbox: Box | EllipsisType) -> MaskedImage:
180 bbox, _ = self._handle_getitem_args(bbox)
181 return self._transfer_metadata(
182 MaskedImage(
183 # Projection and obs_info propagate from the image.
184 self.image[bbox],
185 mask=self.mask[bbox],
186 variance=self.variance[bbox],
187 ),
188 bbox=bbox,
189 )
191 def __setitem__(self, bbox: Box | EllipsisType, value: MaskedImage) -> None:
192 self._image[bbox] = value.image
193 self._mask[bbox] = value.mask
194 self._variance[bbox] = value.variance
196 def __str__(self) -> str:
197 return f"MaskedImage({self.image!s}, {list(self.mask.schema.names)})"
199 def __repr__(self) -> str:
200 return f"MaskedImage({self.image!r}, mask_schema={self.mask.schema!r})"
202 def copy(self) -> MaskedImage:
203 """Deep-copy the masked image and metadata."""
204 return self._transfer_metadata(
205 MaskedImage(image=self._image.copy(), mask=self._mask.copy(), variance=self._variance.copy()),
206 copy=True,
207 )
209 def serialize(self, archive: OutputArchive[Any]) -> MaskedImageSerializationModel[Any]:
210 """Serialize the masked image to an output archive.
212 Parameters
213 ----------
214 archive
215 Archive to write to.
216 """
217 return self._serialize_impl(MaskedImageSerializationModel, archive)
219 def _serialize_impl[M: MaskedImageSerializationModel[Any]](
220 self, model_type: type[M], archive: OutputArchive[Any]
221 ) -> M:
222 serialized_image = archive.serialize_direct(
223 "image", functools.partial(self.image.serialize, save_projection=False)
224 )
225 serialized_mask = archive.serialize_direct(
226 "mask", functools.partial(self.mask.serialize, save_projection=False)
227 )
228 serialized_variance = archive.serialize_direct(
229 "variance", functools.partial(self.variance.serialize, save_projection=False)
230 )
231 serialized_projection = (
232 archive.serialize_direct("sky_projection", self.sky_projection.serialize)
233 if self.sky_projection is not None
234 else None
235 )
236 # When M is a subclass of MaskedImageSerializationModel, it probably
237 # has fields that aren't being set here. We're intentionally making use
238 # of the fact that model_construct doesn't guard against that so we can
239 # instead set them in a subclass implementation later, after calling
240 # super() to construct an instance of the right type. MyPy is actually
241 # fine with this, but only because it incorrectly thinks model_type is
242 # only ever MaskedImageSerializationModel, not a subclass, and that
243 # makes it incorrectly unhappy about the return type.
244 return model_type.model_construct( # type: ignore[return-value]
245 image=serialized_image,
246 mask=serialized_mask,
247 variance=serialized_variance,
248 sky_projection=serialized_projection,
249 metadata=self.metadata,
250 )
252 @staticmethod
253 def _get_archive_tree_type[P: pydantic.BaseModel](
254 pointer_type: type[P],
255 ) -> type[MaskedImageSerializationModel[P]]:
256 """Return the serialization model type for this object for an archive
257 type that uses the given pointer type.
258 """
259 return MaskedImageSerializationModel[pointer_type] # type: ignore
261 @staticmethod
262 def from_legacy(
263 legacy: LegacyMaskedImage,
264 *,
265 unit: astropy.units.UnitBase | None = None,
266 plane_map: Mapping[str, MaskPlane] | None = None,
267 ) -> MaskedImage:
268 """Convert from an `lsst.afw.image.MaskedImage` instance.
270 Parameters
271 ----------
272 legacy
273 An `lsst.afw.image.MaskedImage` instance that will share image and
274 variance (but not mask) pixel data with the returned object.
275 unit
276 Units of the image.
277 plane_map
278 A mapping from legacy mask plane name to the new plane name and
279 description. If not provided, the right legacy mask plane will be
280 guessed, but this can depend on which mask planes the legacy
281 mask actually has set.
282 """
283 return MaskedImage(
284 image=Image.from_legacy(legacy.getImage(), unit),
285 mask=Mask.from_legacy(legacy.getMask(), plane_map),
286 variance=Image.from_legacy(legacy.getVariance()),
287 )
289 def to_legacy(
290 self, *, copy: bool | None = None, plane_map: Mapping[str, MaskPlane] | None = None
291 ) -> LegacyMaskedImage:
292 """Convert to an `lsst.afw.image.MaskedImage` instance.
294 Parameters
295 ----------
296 copy
297 If `True`, always copy the image and variance pixel data.
298 If `False`, return a view, and raise `TypeError` if the pixel data
299 is read-only (this is not supported by afw). If `None`, only copy
300 if the pixel data is read-only. Mask pixel data is always copied.
301 plane_map
302 A mapping from legacy mask plane name to the new plane name and
303 description.
304 """
305 import lsst.afw.image
307 return lsst.afw.image.MaskedImage(
308 self.image.to_legacy(copy=copy),
309 mask=self.mask.to_legacy(plane_map),
310 variance=self.variance.to_legacy(copy=copy),
311 dtype=self.image.array.dtype,
312 )
314 @classmethod
315 def from_hdu_list(
316 cls,
317 hdu_list: astropy.io.fits.HDUList,
318 *,
319 fits_wcs_frame: Frame | None = DEFAULT_PIXEL_FRAME,
320 ) -> MaskedImage:
321 """Reconstruct a `~lsst.images.MaskedImage` from a cut-down
322 ``lsst.images`` FITS HDU list.
324 This assumes the ``PRIMARY``, ``IMAGE``, ``MASK``, and ``VARIANCE``
325 HDUs written for the masked-image cut-outs produced by
326 ``dax_images_cutout``: a real ``lsst.images`` file with its JSON-tree,
327 index, and any nested-archive HDUs dropped. The reconstructed object
328 can be re-serialized as a normal ``lsst.images`` file (with schema and
329 index) so it can be read with the full ``lsst.images`` infrastructure.
331 Parameters
332 ----------
333 hdu_list
334 HDU list with ``IMAGE``, ``MASK``, and ``VARIANCE`` extensions and
335 a primary HDU.
336 fits_wcs_frame
337 Pixel-grid `~lsst.images.Frame` for the
338 `~lsst.images.SkyProjection` reconstructed from the FITS WCS.
339 Defaults to a plain pixel frame; pass `None` to skip attaching a
340 projection.
342 Returns
343 -------
344 `~lsst.images.MaskedImage`
345 The reconstructed masked image.
347 Raises
348 ------
349 ValueError
350 Raised if the ``MASK`` HDU has neither ``MSKN`` nor ``MP_`` mask-
351 plane cards, since the mask schema cannot then be reconstructed, or
352 if ``hdu_list`` contains more than one ``MASK`` HDU (multiple
353 ``MASK`` extensions, distinguished by ``EXTVER``, are not handled
354 here and would otherwise be silently dropped).
356 Notes
357 -----
358 Both mask-plane conventions are supported: the self-describing
359 ``MSKN``/``MSKM``/``MSKD`` cards written by ``lsst.images``, and the
360 legacy `lsst.afw.image` ``MP_*`` cards (as produced by
361 ``dax_images_cutout`` from afw-written images). Legacy masks are
362 mapped to a new schema with the same plane-guessing used by
363 `read_legacy`.
365 Unlike `read_legacy`, the legacy ``MP_*`` mask-plane cards are kept
366 (not stripped) for backwards compatibility, since this path
367 reconstructs a file that may still be read by legacy tooling. They are
368 re-indexed to the reshuffled schema so each ``MP_`` bit matches the
369 plane's position in the written ``MSKN`` layout.
371 The headers of the HDUs in ``hdu_list`` are modified in place: the WCS
372 and mask-schema cards interpreted here are stripped from the caller's
373 headers.
374 """
375 n_mask_hdus = sum(1 for hdu in hdu_list if hdu.name == "MASK")
376 if n_mask_hdus > 1:
377 raise ValueError(
378 f"Found {n_mask_hdus} MASK HDUs; from_hdu_list supports only a single MASK "
379 "extension and would otherwise silently drop mask information from the others."
380 )
381 mask_hdu = hdu_list["MASK"]
382 if not any(card.keyword.startswith(("MSKN", "MP_")) for card in mask_hdu.header.cards):
383 raise ValueError("MASK HDU has no MSKN or MP_ cards; cannot reconstruct the mask schema.")
384 opaque_metadata = fits.FitsOpaqueMetadata()
385 opaque_metadata.add_cutdown_primary_header(hdu_list[0].header)
386 image = Image._read_legacy_hdu(
387 hdu_list["IMAGE"], opaque_metadata, preserve_bintable=None, fits_wcs_frame=fits_wcs_frame
388 )
389 mask = Mask._read_legacy_hdu(
390 mask_hdu, opaque_metadata, fits_wcs_frame=None, strip_legacy_planes=False
391 )
392 variance = Image._read_legacy_hdu(
393 hdu_list["VARIANCE"], opaque_metadata, preserve_bintable=None, fits_wcs_frame=None
394 )
395 result = cls(image, mask=mask, variance=variance)
396 result._opaque_metadata = opaque_metadata
397 return result
399 @staticmethod
400 def read_legacy(
401 uri: ResourcePathExpression,
402 *,
403 preserve_quantization: bool = False,
404 plane_map: Mapping[str, MaskPlane] | None = None,
405 component: Literal["image", "mask", "variance"] | None = None,
406 fits_wcs_frame: Frame | None = None,
407 ) -> Any:
408 """Read a FITS file written by `lsst.afw.image.MaskedImage.writeFits`.
410 Parameters
411 ----------
412 uri
413 URI or file name.
414 preserve_quantization
415 If `True`, ensure that writing the masked image back out again will
416 exactly preserve quantization-compressed pixel values. This causes
417 the image and variance plane arrays to be marked as read-only and
418 stores the original binary table data for those planes in memory.
419 If the `~lsst.images.MaskedImage` is copied, the precompressed
420 pixel values are not transferred to the copy.
421 plane_map
422 A mapping from legacy mask plane name to the new plane name and
423 description. If not provided, the right legacy mask plane will be
424 guessed, but this can depend on which mask planes the legacy
425 mask actually has set.
426 component
427 A component to read instead of the full image.
428 fits_wcs_frame
429 If not `None` and the HDU containing the image plane has a FITS
430 WCS, attach a `~lsst.images.SkyProjection` to the returned masked
431 image by converting that WCS. When ``component`` is one of
432 ``"image"``, ``"mask"``, or ``"variance"``, a FITS WCS from the
433 component HDU is used instead (all three should have the same WCS).
434 """
435 fs, fspath = ResourcePath(uri).to_fsspec()
436 with fs.open(fspath) as stream, astropy.io.fits.open(stream) as hdu_list:
437 return MaskedImage._read_legacy_hdus(
438 hdu_list,
439 uri,
440 preserve_quantization=preserve_quantization,
441 plane_map=plane_map,
442 component=component,
443 fits_wcs_frame=fits_wcs_frame,
444 )
446 @staticmethod
447 def _read_legacy_hdus(
448 hdu_list: astropy.io.fits.HDUList,
449 uri: ResourcePathExpression,
450 *,
451 opaque_metadata: fits.FitsOpaqueMetadata | None = None,
452 preserve_quantization: bool = False,
453 plane_map: Mapping[str, MaskPlane] | None = None,
454 component: Literal["image", "mask", "variance"] | None,
455 fits_wcs_frame: Frame | None = None,
456 ) -> Any:
457 if opaque_metadata is None:
458 opaque_metadata = fits.FitsOpaqueMetadata()
459 opaque_metadata.extract_legacy_primary_header(hdu_list[0].header)
460 image_bintable_hdu: astropy.io.fits.BinTableHDU | None = None
461 variance_bintable_hdu: astropy.io.fits.BinTableHDU | None = None
462 result: Any
463 with ExitStack() as exit_stack:
464 if preserve_quantization:
465 fs, fspath = ResourcePath(uri).to_fsspec()
466 bintable_stream = exit_stack.enter_context(fs.open(fspath))
467 bintable_hdu_list = exit_stack.enter_context(
468 astropy.io.fits.open(bintable_stream, disable_image_compression=True)
469 )
470 image_bintable_hdu = bintable_hdu_list[1]
471 variance_bintable_hdu = bintable_hdu_list[3]
472 if component is None or component == "image":
473 image = Image._read_legacy_hdu(
474 hdu_list[1],
475 opaque_metadata,
476 preserve_bintable=image_bintable_hdu,
477 fits_wcs_frame=fits_wcs_frame,
478 )
479 if component == "image":
480 result = image
481 if component is None or component == "mask":
482 mask = Mask._read_legacy_hdu(
483 hdu_list[2],
484 opaque_metadata,
485 plane_map=plane_map,
486 fits_wcs_frame=fits_wcs_frame if component is not None else None,
487 )
488 if component == "mask":
489 result = mask
490 if component is None or component == "variance":
491 variance = Image._read_legacy_hdu(
492 hdu_list[3],
493 opaque_metadata,
494 preserve_bintable=variance_bintable_hdu,
495 fits_wcs_frame=fits_wcs_frame if component is not None else None,
496 )
497 if component == "variance":
498 result = variance
499 if component is None:
500 result = MaskedImage(image, mask=mask, variance=variance)
501 result._opaque_metadata = opaque_metadata
502 return result
504 def _fill_legacy_metadata(self, legacy_metadata: MutableMapping[str, Any]) -> None:
505 """Fill a legacy mutable mapping (e.g `lsst.daf.base.PropertySet`)
506 with metadata suitable for an `lsst.afw.image.Exposure` representation
507 of this object.
508 """
509 # We just dump all of the FITS headers and non-FITS metadata into the
510 # legacy metadata component, to make sure we have everything. We dump
511 # the latter into a pair of special cards to be able to full round-trip
512 # them (including case preservation).
513 if self.unit is not None:
514 try:
515 legacy_metadata["BUNIT"] = self.unit.to_string(format="fits")
516 except ValueError:
517 # Write units that astropy doesn't think FITS will accept
518 # anyway; FITS standard says "SHOULD" about using its
519 # recommended units, and coloring outside the lines is better
520 # than lying.
521 legacy_metadata["BUNIT"] = self.unit.to_string()
522 if isinstance(self._opaque_metadata, fits.FitsOpaqueMetadata):
523 legacy_metadata.update(self._opaque_metadata.headers[fits.ExtensionKey()])
524 for n, (k, v) in enumerate(self.metadata.items()):
525 legacy_metadata[f"LSST IMAGES KEY {n + 1}"] = k
526 legacy_metadata[f"LSST IMAGES VALUE {n + 1}"] = v
529class MaskedImageSerializationModel[P: pydantic.BaseModel](ArchiveTree):
530 """A Pydantic model used to represent a serialized `MaskedImage`."""
532 SCHEMA_NAME: ClassVar[str] = "masked_image"
533 SCHEMA_VERSION: ClassVar[str] = "1.0.0"
534 MIN_READ_VERSION: ClassVar[int] = 1
535 PUBLIC_TYPE: ClassVar[type] = MaskedImage
537 image: ImageSerializationModel[P] = pydantic.Field(description="The main data image.")
538 mask: MaskSerializationModel[P] = pydantic.Field(
539 description="Bitmask that annotates the main image's pixels."
540 )
541 variance: ImageSerializationModel[P] = pydantic.Field(
542 description="Per-pixel variance estimates for the main image."
543 )
544 sky_projection: SkyProjectionSerializationModel[P] | None = pydantic.Field(
545 default=None,
546 exclude_if=is_none,
547 description="Projection that maps the pixel grid to the sky.",
548 )
550 @property
551 def bbox(self) -> Box:
552 """The bounding box of the image."""
553 return self.image.bbox
555 def deserialize(
556 self, archive: InputArchive[Any], *, bbox: Box | None = None, **kwargs: Any
557 ) -> MaskedImage:
558 """Deserialize an image from an input archive.
560 Parameters
561 ----------
562 archive
563 Archive to read from.
564 bbox
565 Bounding box of a subimage to read instead.
566 **kwargs
567 Unsupported keyword arguments are accepted only to provide better
568 error messages (raising `serialization.InvalidParameterError`).
569 """
570 if kwargs: 570 ↛ 571line 570 didn't jump to line 571 because the condition on line 570 was never true
571 raise InvalidParameterError(f"Unrecognized parameters for MaskedImage: {set(kwargs.keys())}.")
572 image = self.image.deserialize(archive, bbox=bbox)
573 mask = self.mask.deserialize(archive, bbox=bbox)
574 variance = self.variance.deserialize(archive, bbox=bbox)
575 sky_projection = self.sky_projection.deserialize(archive) if self.sky_projection is not None else None
576 return MaskedImage(
577 image, mask=mask, variance=variance, sky_projection=sky_projection
578 )._finish_deserialize(self)
580 def deserialize_component(self, component: str, archive: InputArchive[Any], **kwargs: Any) -> Any:
581 if component == "bbox" and kwargs: 581 ↛ 582line 581 didn't jump to line 582 because the condition on line 581 was never true
582 raise InvalidParameterError(
583 f"Unrecognized parameters for MaskedImage.bbox: {set(kwargs.keys())}."
584 )
585 return super().deserialize_component(component, archive, **kwargs)