Coverage for python/lsst/images/_masked_image.py: 70%
194 statements
« prev ^ index » next coverage.py v7.15.0, created at 2026-07-04 17:19 +0000
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 if bbox is ...:
181 return self
182 super().__getitem__(bbox)
183 return self._transfer_metadata(
184 MaskedImage(
185 # Projection and obs_info propagate from the image.
186 self.image[bbox],
187 mask=self.mask[bbox],
188 variance=self.variance[bbox],
189 ),
190 bbox=bbox,
191 )
193 def __setitem__(self, bbox: Box | EllipsisType, value: MaskedImage) -> None:
194 self._image[bbox] = value.image
195 self._mask[bbox] = value.mask
196 self._variance[bbox] = value.variance
198 def __str__(self) -> str:
199 return f"MaskedImage({self.image!s}, {list(self.mask.schema.names)})"
201 def __repr__(self) -> str:
202 return f"MaskedImage({self.image!r}, mask_schema={self.mask.schema!r})"
204 def copy(self) -> MaskedImage:
205 """Deep-copy the masked image and metadata."""
206 return self._transfer_metadata(
207 MaskedImage(image=self._image.copy(), mask=self._mask.copy(), variance=self._variance.copy()),
208 copy=True,
209 )
211 def serialize(self, archive: OutputArchive[Any]) -> MaskedImageSerializationModel[Any]:
212 """Serialize the masked image to an output archive.
214 Parameters
215 ----------
216 archive
217 Archive to write to.
218 """
219 return self._serialize_impl(MaskedImageSerializationModel, archive)
221 def _serialize_impl[M: MaskedImageSerializationModel[Any]](
222 self, model_type: type[M], archive: OutputArchive[Any]
223 ) -> M:
224 serialized_image = archive.serialize_direct(
225 "image", functools.partial(self.image.serialize, save_projection=False)
226 )
227 serialized_mask = archive.serialize_direct(
228 "mask", functools.partial(self.mask.serialize, save_projection=False)
229 )
230 serialized_variance = archive.serialize_direct(
231 "variance", functools.partial(self.variance.serialize, save_projection=False)
232 )
233 serialized_projection = (
234 archive.serialize_direct("sky_projection", self.sky_projection.serialize)
235 if self.sky_projection is not None
236 else None
237 )
238 # When M is a subclass of MaskedImageSerializationModel, it probably
239 # has fields that aren't being set here. We're intentionally making use
240 # of the fact that model_construct doesn't guard against that so we can
241 # instead set them in a subclass implementation later, after calling
242 # super() to construct an instance of the right type. MyPy is actually
243 # fine with this, but only because it incorrectly thinks model_type is
244 # only ever MaskedImageSerializationModel, not a subclass, and that
245 # makes it incorrectly unhappy about the return type.
246 return model_type.model_construct( # type: ignore[return-value]
247 image=serialized_image,
248 mask=serialized_mask,
249 variance=serialized_variance,
250 sky_projection=serialized_projection,
251 metadata=self.metadata,
252 )
254 @staticmethod
255 def _get_archive_tree_type[P: pydantic.BaseModel](
256 pointer_type: type[P],
257 ) -> type[MaskedImageSerializationModel[P]]:
258 """Return the serialization model type for this object for an archive
259 type that uses the given pointer type.
260 """
261 return MaskedImageSerializationModel[pointer_type] # type: ignore
263 @staticmethod
264 def from_legacy(
265 legacy: LegacyMaskedImage,
266 *,
267 unit: astropy.units.UnitBase | None = None,
268 plane_map: Mapping[str, MaskPlane] | None = None,
269 ) -> MaskedImage:
270 """Convert from an `lsst.afw.image.MaskedImage` instance.
272 Parameters
273 ----------
274 legacy
275 An `lsst.afw.image.MaskedImage` instance that will share image and
276 variance (but not mask) pixel data with the returned object.
277 unit
278 Units of the image.
279 plane_map
280 A mapping from legacy mask plane name to the new plane name and
281 description. If not provided, the right legacy mask plane will be
282 guessed, but this can depend on which mask planes the legacy
283 mask actually has set.
284 """
285 return MaskedImage(
286 image=Image.from_legacy(legacy.getImage(), unit),
287 mask=Mask.from_legacy(legacy.getMask(), plane_map),
288 variance=Image.from_legacy(legacy.getVariance()),
289 )
291 def to_legacy(
292 self, *, copy: bool | None = None, plane_map: Mapping[str, MaskPlane] | None = None
293 ) -> LegacyMaskedImage:
294 """Convert to an `lsst.afw.image.MaskedImage` instance.
296 Parameters
297 ----------
298 copy
299 If `True`, always copy the image and variance pixel data.
300 If `False`, return a view, and raise `TypeError` if the pixel data
301 is read-only (this is not supported by afw). If `None`, only copy
302 if the pixel data is read-only. Mask pixel data is always copied.
303 plane_map
304 A mapping from legacy mask plane name to the new plane name and
305 description.
306 """
307 import lsst.afw.image
309 return lsst.afw.image.MaskedImage(
310 self.image.to_legacy(copy=copy),
311 mask=self.mask.to_legacy(plane_map),
312 variance=self.variance.to_legacy(copy=copy),
313 dtype=self.image.array.dtype,
314 )
316 @classmethod
317 def from_hdu_list(
318 cls,
319 hdu_list: astropy.io.fits.HDUList,
320 *,
321 fits_wcs_frame: Frame | None = DEFAULT_PIXEL_FRAME,
322 ) -> MaskedImage:
323 """Reconstruct a `~lsst.images.MaskedImage` from a cut-down
324 ``lsst.images`` FITS HDU list.
326 This assumes the ``PRIMARY``, ``IMAGE``, ``MASK``, and ``VARIANCE``
327 HDUs written for the masked-image cut-outs produced by
328 ``dax_images_cutout``: a real ``lsst.images`` file with its JSON-tree,
329 index, and any nested-archive HDUs dropped. The reconstructed object
330 can be re-serialized as a normal ``lsst.images`` file (with schema and
331 index) so it can be read with the full ``lsst.images`` infrastructure.
333 Parameters
334 ----------
335 hdu_list
336 HDU list with ``IMAGE``, ``MASK``, and ``VARIANCE`` extensions and
337 a primary HDU.
338 fits_wcs_frame
339 Pixel-grid `~lsst.images.Frame` for the
340 `~lsst.images.SkyProjection` reconstructed from the FITS WCS.
341 Defaults to a plain pixel frame; pass `None` to skip attaching a
342 projection.
344 Returns
345 -------
346 `~lsst.images.MaskedImage`
347 The reconstructed masked image.
349 Raises
350 ------
351 ValueError
352 Raised if the ``MASK`` HDU has neither ``MSKN`` nor ``MP_`` mask-
353 plane cards, since the mask schema cannot then be reconstructed, or
354 if ``hdu_list`` contains more than one ``MASK`` HDU (multiple
355 ``MASK`` extensions, distinguished by ``EXTVER``, are not handled
356 here and would otherwise be silently dropped).
358 Notes
359 -----
360 Both mask-plane conventions are supported: the self-describing
361 ``MSKN``/``MSKM``/``MSKD`` cards written by ``lsst.images``, and the
362 legacy `lsst.afw.image` ``MP_*`` cards (as produced by
363 ``dax_images_cutout`` from afw-written images). Legacy masks are
364 mapped to a new schema with the same plane-guessing used by
365 `read_legacy`.
367 Unlike `read_legacy`, the legacy ``MP_*`` mask-plane cards are kept
368 (not stripped) for backwards compatibility, since this path
369 reconstructs a file that may still be read by legacy tooling. They are
370 re-indexed to the reshuffled schema so each ``MP_`` bit matches the
371 plane's position in the written ``MSKN`` layout.
373 The headers of the HDUs in ``hdu_list`` are modified in place: the WCS
374 and mask-schema cards interpreted here are stripped from the caller's
375 headers.
376 """
377 n_mask_hdus = sum(1 for hdu in hdu_list if hdu.name == "MASK")
378 if n_mask_hdus > 1:
379 raise ValueError(
380 f"Found {n_mask_hdus} MASK HDUs; from_hdu_list supports only a single MASK "
381 "extension and would otherwise silently drop mask information from the others."
382 )
383 mask_hdu = hdu_list["MASK"]
384 if not any(card.keyword.startswith(("MSKN", "MP_")) for card in mask_hdu.header.cards):
385 raise ValueError("MASK HDU has no MSKN or MP_ cards; cannot reconstruct the mask schema.")
386 opaque_metadata = fits.FitsOpaqueMetadata()
387 opaque_metadata.add_cutdown_primary_header(hdu_list[0].header)
388 image = Image._read_legacy_hdu(
389 hdu_list["IMAGE"], opaque_metadata, preserve_bintable=None, fits_wcs_frame=fits_wcs_frame
390 )
391 mask = Mask._read_legacy_hdu(
392 mask_hdu, opaque_metadata, fits_wcs_frame=None, strip_legacy_planes=False
393 )
394 variance = Image._read_legacy_hdu(
395 hdu_list["VARIANCE"], opaque_metadata, preserve_bintable=None, fits_wcs_frame=None
396 )
397 result = cls(image, mask=mask, variance=variance)
398 result._opaque_metadata = opaque_metadata
399 return result
401 @staticmethod
402 def read_legacy(
403 uri: ResourcePathExpression,
404 *,
405 preserve_quantization: bool = False,
406 plane_map: Mapping[str, MaskPlane] | None = None,
407 component: Literal["image", "mask", "variance"] | None = None,
408 fits_wcs_frame: Frame | None = None,
409 ) -> Any:
410 """Read a FITS file written by `lsst.afw.image.MaskedImage.writeFits`.
412 Parameters
413 ----------
414 uri
415 URI or file name.
416 preserve_quantization
417 If `True`, ensure that writing the masked image back out again will
418 exactly preserve quantization-compressed pixel values. This causes
419 the image and variance plane arrays to be marked as read-only and
420 stores the original binary table data for those planes in memory.
421 If the `~lsst.images.MaskedImage` is copied, the precompressed
422 pixel values are not transferred to the copy.
423 plane_map
424 A mapping from legacy mask plane name to the new plane name and
425 description. If not provided, the right legacy mask plane will be
426 guessed, but this can depend on which mask planes the legacy
427 mask actually has set.
428 component
429 A component to read instead of the full image.
430 fits_wcs_frame
431 If not `None` and the HDU containing the image plane has a FITS
432 WCS, attach a `~lsst.images.SkyProjection` to the returned masked
433 image by converting that WCS. When ``component`` is one of
434 ``"image"``, ``"mask"``, or ``"variance"``, a FITS WCS from the
435 component HDU is used instead (all three should have the same WCS).
436 """
437 fs, fspath = ResourcePath(uri).to_fsspec()
438 with fs.open(fspath) as stream, astropy.io.fits.open(stream) as hdu_list:
439 return MaskedImage._read_legacy_hdus(
440 hdu_list,
441 uri,
442 preserve_quantization=preserve_quantization,
443 plane_map=plane_map,
444 component=component,
445 fits_wcs_frame=fits_wcs_frame,
446 )
448 @staticmethod
449 def _read_legacy_hdus(
450 hdu_list: astropy.io.fits.HDUList,
451 uri: ResourcePathExpression,
452 *,
453 opaque_metadata: fits.FitsOpaqueMetadata | None = None,
454 preserve_quantization: bool = False,
455 plane_map: Mapping[str, MaskPlane] | None = None,
456 component: Literal["image", "mask", "variance"] | None,
457 fits_wcs_frame: Frame | None = None,
458 ) -> Any:
459 if opaque_metadata is None:
460 opaque_metadata = fits.FitsOpaqueMetadata()
461 opaque_metadata.extract_legacy_primary_header(hdu_list[0].header)
462 image_bintable_hdu: astropy.io.fits.BinTableHDU | None = None
463 variance_bintable_hdu: astropy.io.fits.BinTableHDU | None = None
464 result: Any
465 with ExitStack() as exit_stack:
466 if preserve_quantization:
467 fs, fspath = ResourcePath(uri).to_fsspec()
468 bintable_stream = exit_stack.enter_context(fs.open(fspath))
469 bintable_hdu_list = exit_stack.enter_context(
470 astropy.io.fits.open(bintable_stream, disable_image_compression=True)
471 )
472 image_bintable_hdu = bintable_hdu_list[1]
473 variance_bintable_hdu = bintable_hdu_list[3]
474 if component is None or component == "image":
475 image = Image._read_legacy_hdu(
476 hdu_list[1],
477 opaque_metadata,
478 preserve_bintable=image_bintable_hdu,
479 fits_wcs_frame=fits_wcs_frame,
480 )
481 if component == "image":
482 result = image
483 if component is None or component == "mask":
484 mask = Mask._read_legacy_hdu(
485 hdu_list[2],
486 opaque_metadata,
487 plane_map=plane_map,
488 fits_wcs_frame=fits_wcs_frame if component is not None else None,
489 )
490 if component == "mask":
491 result = mask
492 if component is None or component == "variance":
493 variance = Image._read_legacy_hdu(
494 hdu_list[3],
495 opaque_metadata,
496 preserve_bintable=variance_bintable_hdu,
497 fits_wcs_frame=fits_wcs_frame if component is not None else None,
498 )
499 if component == "variance":
500 result = variance
501 if component is None:
502 result = MaskedImage(image, mask=mask, variance=variance)
503 result._opaque_metadata = opaque_metadata
504 return result
506 def _fill_legacy_metadata(self, legacy_metadata: MutableMapping[str, Any]) -> None:
507 """Fill a legacy mutable mapping (e.g `lsst.daf.base.PropertySet`)
508 with metadata suitable for an `lsst.afw.image.Exposure` representation
509 of this object.
510 """
511 # We just dump all of the FITS headers and non-FITS metadata into the
512 # legacy metadata component, to make sure we have everything. We dump
513 # the latter into a pair of special cards to be able to full round-trip
514 # them (including case preservation).
515 if self.unit is not None:
516 try:
517 legacy_metadata["BUNIT"] = self.unit.to_string(format="fits")
518 except ValueError:
519 # Write units that astropy doesn't think FITS will accept
520 # anyway; FITS standard says "SHOULD" about using its
521 # recommended units, and coloring outside the lines is better
522 # than lying.
523 legacy_metadata["BUNIT"] = self.unit.to_string()
524 if isinstance(self._opaque_metadata, fits.FitsOpaqueMetadata):
525 legacy_metadata.update(self._opaque_metadata.headers[fits.ExtensionKey()])
526 for n, (k, v) in enumerate(self.metadata.items()):
527 legacy_metadata[f"LSST IMAGES KEY {n + 1}"] = k
528 legacy_metadata[f"LSST IMAGES VALUE {n + 1}"] = v
531class MaskedImageSerializationModel[P: pydantic.BaseModel](ArchiveTree):
532 """A Pydantic model used to represent a serialized `MaskedImage`."""
534 SCHEMA_NAME: ClassVar[str] = "masked_image"
535 SCHEMA_VERSION: ClassVar[str] = "1.0.0"
536 MIN_READ_VERSION: ClassVar[int] = 1
537 PUBLIC_TYPE: ClassVar[type] = MaskedImage
539 image: ImageSerializationModel[P] = pydantic.Field(description="The main data image.")
540 mask: MaskSerializationModel[P] = pydantic.Field(
541 description="Bitmask that annotates the main image's pixels."
542 )
543 variance: ImageSerializationModel[P] = pydantic.Field(
544 description="Per-pixel variance estimates for the main image."
545 )
546 sky_projection: SkyProjectionSerializationModel[P] | None = pydantic.Field(
547 default=None,
548 exclude_if=is_none,
549 description="Projection that maps the pixel grid to the sky.",
550 )
552 @property
553 def bbox(self) -> Box:
554 """The bounding box of the image."""
555 return self.image.bbox
557 def deserialize(
558 self, archive: InputArchive[Any], *, bbox: Box | None = None, **kwargs: Any
559 ) -> MaskedImage:
560 """Deserialize an image from an input archive.
562 Parameters
563 ----------
564 archive
565 Archive to read from.
566 bbox
567 Bounding box of a subimage to read instead.
568 **kwargs
569 Unsupported keyword arguments are accepted only to provide better
570 error messages (raising `serialization.InvalidParameterError`).
571 """
572 if kwargs: 572 ↛ 573line 572 didn't jump to line 573 because the condition on line 572 was never true
573 raise InvalidParameterError(f"Unrecognized parameters for MaskedImage: {set(kwargs.keys())}.")
574 image = self.image.deserialize(archive, bbox=bbox)
575 mask = self.mask.deserialize(archive, bbox=bbox)
576 variance = self.variance.deserialize(archive, bbox=bbox)
577 sky_projection = self.sky_projection.deserialize(archive) if self.sky_projection is not None else None
578 return MaskedImage(
579 image, mask=mask, variance=variance, sky_projection=sky_projection
580 )._finish_deserialize(self)
582 def deserialize_component(self, component: str, archive: InputArchive[Any], **kwargs: Any) -> Any:
583 if component == "bbox" and kwargs: 583 ↛ 584line 583 didn't jump to line 584 because the condition on line 583 was never true
584 raise InvalidParameterError(
585 f"Unrecognized parameters for MaskedImage.bbox: {set(kwargs.keys())}."
586 )
587 return super().deserialize_component(component, archive, **kwargs)