Coverage for python/lsst/images/_masked_image.py: 42%
192 statements
« prev ^ index » next coverage.py v7.15.2, created at 2026-07-24 02:46 +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 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 sky_projection: SkyProjection[Any] | None = None,
268 ) -> MaskedImage:
269 """Convert from an `lsst.afw.image.MaskedImage` instance.
271 Parameters
272 ----------
273 legacy
274 An `lsst.afw.image.MaskedImage` instance that will share image and
275 variance (but not mask) pixel data with the returned object.
276 unit
277 Units of the image.
278 plane_map
279 A mapping from legacy mask plane name to the new plane name and
280 description. If not provided, the right legacy mask plane will be
281 guessed, but this can depend on which mask planes the legacy
282 mask actually has set.
283 sky_projection
284 Projection from pixels to xky.
285 """
286 return MaskedImage(
287 image=Image.from_legacy(legacy.getImage(), unit, sky_projection=sky_projection),
288 mask=Mask.from_legacy(legacy.getMask(), plane_map, sky_projection=sky_projection),
289 variance=Image.from_legacy(legacy.getVariance(), sky_projection=sky_projection),
290 )
292 def to_legacy(
293 self, *, copy: bool | None = None, plane_map: Mapping[str, MaskPlane] | None = None
294 ) -> LegacyMaskedImage:
295 """Convert to an `lsst.afw.image.MaskedImage` instance.
297 Parameters
298 ----------
299 copy
300 If `True`, always copy the image and variance pixel data.
301 If `False`, return a view, and raise `TypeError` if the pixel data
302 is read-only (this is not supported by afw). If `None`, only copy
303 if the pixel data is read-only. Mask pixel data is always copied.
304 plane_map
305 A mapping from legacy mask plane name to the new plane name and
306 description.
307 """
308 import lsst.afw.image
310 return lsst.afw.image.MaskedImage(
311 self.image.to_legacy(copy=copy),
312 mask=self.mask.to_legacy(plane_map),
313 variance=self.variance.to_legacy(copy=copy),
314 dtype=self.image.array.dtype,
315 )
317 @classmethod
318 def from_hdu_list(
319 cls,
320 hdu_list: astropy.io.fits.HDUList,
321 *,
322 fits_wcs_frame: Frame | None = DEFAULT_PIXEL_FRAME,
323 ) -> MaskedImage:
324 """Reconstruct a `~lsst.images.MaskedImage` from a cut-down
325 ``lsst.images`` FITS HDU list.
327 This assumes the ``PRIMARY``, ``IMAGE``, ``MASK``, and ``VARIANCE``
328 HDUs written for the masked-image cut-outs produced by
329 ``dax_images_cutout``: a real ``lsst.images`` file with its JSON-tree,
330 index, and any nested-archive HDUs dropped. The reconstructed object
331 can be re-serialized as a normal ``lsst.images`` file (with schema and
332 index) so it can be read with the full ``lsst.images`` infrastructure.
334 Parameters
335 ----------
336 hdu_list
337 HDU list with ``IMAGE``, ``MASK``, and ``VARIANCE`` extensions and
338 a primary HDU.
339 fits_wcs_frame
340 Pixel-grid `~lsst.images.Frame` for the
341 `~lsst.images.SkyProjection` reconstructed from the FITS WCS.
342 Defaults to a plain pixel frame; pass `None` to skip attaching a
343 projection.
345 Returns
346 -------
347 `~lsst.images.MaskedImage`
348 The reconstructed masked image.
350 Raises
351 ------
352 ValueError
353 Raised if the ``MASK`` HDU has neither ``MSKN`` nor ``MP_`` mask-
354 plane cards, since the mask schema cannot then be reconstructed, or
355 if ``hdu_list`` contains more than one ``MASK`` HDU (multiple
356 ``MASK`` extensions, distinguished by ``EXTVER``, are not handled
357 here and would otherwise be silently dropped).
359 Notes
360 -----
361 Both mask-plane conventions are supported: the self-describing
362 ``MSKN``/``MSKM``/``MSKD`` cards written by ``lsst.images``, and the
363 legacy `lsst.afw.image` ``MP_*`` cards (as produced by
364 ``dax_images_cutout`` from afw-written images). Legacy masks are
365 mapped to a new schema with the same plane-guessing used by
366 `read_legacy`.
368 Unlike `read_legacy`, the legacy ``MP_*`` mask-plane cards are kept
369 (not stripped) for backwards compatibility, since this path
370 reconstructs a file that may still be read by legacy tooling. They are
371 re-indexed to the reshuffled schema so each ``MP_`` bit matches the
372 plane's position in the written ``MSKN`` layout.
374 The headers of the HDUs in ``hdu_list`` are modified in place: the WCS
375 and mask-schema cards interpreted here are stripped from the caller's
376 headers.
377 """
378 n_mask_hdus = sum(1 for hdu in hdu_list if hdu.name == "MASK")
379 if n_mask_hdus > 1:
380 raise ValueError(
381 f"Found {n_mask_hdus} MASK HDUs; from_hdu_list supports only a single MASK "
382 "extension and would otherwise silently drop mask information from the others."
383 )
384 mask_hdu = hdu_list["MASK"]
385 if not any(card.keyword.startswith(("MSKN", "MP_")) for card in mask_hdu.header.cards):
386 raise ValueError("MASK HDU has no MSKN or MP_ cards; cannot reconstruct the mask schema.")
387 opaque_metadata = fits.FitsOpaqueMetadata()
388 opaque_metadata.add_cutdown_primary_header(hdu_list[0].header)
389 image = Image._read_legacy_hdu(
390 hdu_list["IMAGE"], opaque_metadata, preserve_bintable=None, fits_wcs_frame=fits_wcs_frame
391 )
392 mask = Mask._read_legacy_hdu(
393 mask_hdu, opaque_metadata, fits_wcs_frame=None, strip_legacy_planes=False
394 )
395 variance = Image._read_legacy_hdu(
396 hdu_list["VARIANCE"], opaque_metadata, preserve_bintable=None, fits_wcs_frame=None
397 )
398 result = cls(image, mask=mask, variance=variance)
399 result._opaque_metadata = opaque_metadata
400 return result
402 @staticmethod
403 def read_legacy(
404 uri: ResourcePathExpression,
405 *,
406 preserve_quantization: bool = False,
407 plane_map: Mapping[str, MaskPlane] | None = None,
408 component: Literal["image", "mask", "variance"] | None = None,
409 fits_wcs_frame: Frame | None = None,
410 ) -> Any:
411 """Read a FITS file written by `lsst.afw.image.MaskedImage.writeFits`.
413 Parameters
414 ----------
415 uri
416 URI or file name.
417 preserve_quantization
418 If `True`, ensure that writing the masked image back out again will
419 exactly preserve quantization-compressed pixel values. This causes
420 the image and variance plane arrays to be marked as read-only and
421 stores the original binary table data for those planes in memory.
422 If the `~lsst.images.MaskedImage` is copied, the precompressed
423 pixel values are not transferred to the copy.
424 plane_map
425 A mapping from legacy mask plane name to the new plane name and
426 description. If not provided, the right legacy mask plane will be
427 guessed, but this can depend on which mask planes the legacy
428 mask actually has set.
429 component
430 A component to read instead of the full image.
431 fits_wcs_frame
432 If not `None` and the HDU containing the image plane has a FITS
433 WCS, attach a `~lsst.images.SkyProjection` to the returned masked
434 image by converting that WCS. When ``component`` is one of
435 ``"image"``, ``"mask"``, or ``"variance"``, a FITS WCS from the
436 component HDU is used instead (all three should have the same WCS).
437 """
438 fs, fspath = ResourcePath(uri).to_fsspec()
439 with fs.open(fspath) as stream, astropy.io.fits.open(stream) as hdu_list:
440 return MaskedImage._read_legacy_hdus(
441 hdu_list,
442 uri,
443 preserve_quantization=preserve_quantization,
444 plane_map=plane_map,
445 component=component,
446 fits_wcs_frame=fits_wcs_frame,
447 )
449 @staticmethod
450 def _read_legacy_hdus(
451 hdu_list: astropy.io.fits.HDUList,
452 uri: ResourcePathExpression,
453 *,
454 opaque_metadata: fits.FitsOpaqueMetadata | None = None,
455 preserve_quantization: bool = False,
456 plane_map: Mapping[str, MaskPlane] | None = None,
457 component: Literal["image", "mask", "variance"] | None,
458 fits_wcs_frame: Frame | None = None,
459 ) -> Any:
460 if opaque_metadata is None:
461 opaque_metadata = fits.FitsOpaqueMetadata()
462 opaque_metadata.extract_legacy_primary_header(hdu_list[0].header)
463 image_bintable_hdu: astropy.io.fits.BinTableHDU | None = None
464 variance_bintable_hdu: astropy.io.fits.BinTableHDU | None = None
465 result: Any
466 with ExitStack() as exit_stack:
467 if preserve_quantization:
468 fs, fspath = ResourcePath(uri).to_fsspec()
469 bintable_stream = exit_stack.enter_context(fs.open(fspath))
470 bintable_hdu_list = exit_stack.enter_context(
471 astropy.io.fits.open(bintable_stream, disable_image_compression=True)
472 )
473 image_bintable_hdu = bintable_hdu_list[1]
474 variance_bintable_hdu = bintable_hdu_list[3]
475 if component is None or component == "image":
476 image = Image._read_legacy_hdu(
477 hdu_list[1],
478 opaque_metadata,
479 preserve_bintable=image_bintable_hdu,
480 fits_wcs_frame=fits_wcs_frame,
481 )
482 if component == "image":
483 result = image
484 if component is None or component == "mask":
485 mask = Mask._read_legacy_hdu(
486 hdu_list[2],
487 opaque_metadata,
488 plane_map=plane_map,
489 fits_wcs_frame=fits_wcs_frame if component is not None else None,
490 )
491 if component == "mask":
492 result = mask
493 if component is None or component == "variance":
494 variance = Image._read_legacy_hdu(
495 hdu_list[3],
496 opaque_metadata,
497 preserve_bintable=variance_bintable_hdu,
498 fits_wcs_frame=fits_wcs_frame if component is not None else None,
499 )
500 if component == "variance":
501 result = variance
502 if component is None:
503 result = MaskedImage(image, mask=mask, variance=variance)
504 result._opaque_metadata = opaque_metadata
505 return result
507 def _fill_legacy_metadata(self, legacy_metadata: MutableMapping[str, Any]) -> None:
508 """Fill a legacy mutable mapping (e.g `lsst.daf.base.PropertySet`)
509 with metadata suitable for an `lsst.afw.image.Exposure` representation
510 of this object.
511 """
512 # We just dump all of the FITS headers and non-FITS metadata into the
513 # legacy metadata component, to make sure we have everything. We dump
514 # the latter into a pair of special cards to be able to full round-trip
515 # them (including case preservation).
516 if self.unit is not None:
517 try:
518 legacy_metadata["BUNIT"] = self.unit.to_string(format="fits")
519 except ValueError:
520 # Write units that astropy doesn't think FITS will accept
521 # anyway; FITS standard says "SHOULD" about using its
522 # recommended units, and coloring outside the lines is better
523 # than lying.
524 legacy_metadata["BUNIT"] = self.unit.to_string()
525 if isinstance(self._opaque_metadata, fits.FitsOpaqueMetadata):
526 legacy_metadata.update(self._opaque_metadata.headers[fits.ExtensionKey()])
527 for n, (k, v) in enumerate(self.metadata.items()):
528 legacy_metadata[f"LSST IMAGES KEY {n + 1}"] = k
529 legacy_metadata[f"LSST IMAGES VALUE {n + 1}"] = v
532class MaskedImageSerializationModel[P: pydantic.BaseModel](ArchiveTree):
533 """A Pydantic model used to represent a serialized `MaskedImage`."""
535 SCHEMA_NAME: ClassVar[str] = "masked_image"
536 SCHEMA_VERSION: ClassVar[str] = "1.0.0.dev0"
537 MIN_READ_VERSION: ClassVar[int] = 1
538 PUBLIC_TYPE: ClassVar[type] = MaskedImage
540 image: ImageSerializationModel[P] = pydantic.Field(description="The main data image.")
541 mask: MaskSerializationModel[P] = pydantic.Field(
542 description="Bitmask that annotates the main image's pixels."
543 )
544 variance: ImageSerializationModel[P] = pydantic.Field(
545 description="Per-pixel variance estimates for the main image."
546 )
547 sky_projection: SkyProjectionSerializationModel[P] | None = pydantic.Field(
548 default=None,
549 exclude_if=is_none,
550 description="Projection that maps the pixel grid to the sky.",
551 )
553 @property
554 def bbox(self) -> Box:
555 """The bounding box of the image."""
556 return self.image.bbox
558 def deserialize(
559 self, archive: InputArchive[Any], *, bbox: Box | None = None, **kwargs: Any
560 ) -> MaskedImage:
561 """Deserialize an image from an input archive.
563 Parameters
564 ----------
565 archive
566 Archive to read from.
567 bbox
568 Bounding box of a subimage to read instead.
569 **kwargs
570 Unsupported keyword arguments are accepted only to provide better
571 error messages (raising `serialization.InvalidParameterError`).
572 """
573 if kwargs: 573 ↛ 574line 573 didn't jump to line 574 because the condition on line 573 was never true
574 raise InvalidParameterError(f"Unrecognized parameters for MaskedImage: {set(kwargs.keys())}.")
575 image = self.image.deserialize(archive, bbox=bbox)
576 mask = self.mask.deserialize(archive, bbox=bbox)
577 variance = self.variance.deserialize(archive, bbox=bbox)
578 sky_projection = self.sky_projection.deserialize(archive) if self.sky_projection is not None else None
579 return MaskedImage(
580 image, mask=mask, variance=variance, sky_projection=sky_projection
581 )._finish_deserialize(self)
583 def deserialize_component(self, component: str, archive: InputArchive[Any], **kwargs: Any) -> Any:
584 if component == "bbox" and kwargs: 584 ↛ 585line 584 didn't jump to line 585 because the condition on line 584 was never true
585 raise InvalidParameterError(
586 f"Unrecognized parameters for MaskedImage.bbox: {set(kwargs.keys())}."
587 )
588 return super().deserialize_component(component, archive, **kwargs)