Coverage for python/lsst/images/_masked_image.py: 31%
160 statements
« prev ^ index » next coverage.py v7.14.1, created at 2026-06-06 01:43 -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
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 Image, ImageSerializationModel
34from ._mask import Mask, MaskPlane, MaskSchema, MaskSerializationModel
35from ._transforms import Frame, Projection, ProjectionSerializationModel
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 `Projection`, it will be used
59 for all planes unless a ``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 projection
63 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 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 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 projection: Projection | None = None,
87 metadata: dict[str, MetadataValue] | None = None,
88 ):
89 super().__init__(metadata)
90 if projection is None:
91 projection = image.projection
92 else:
93 image = image.view(projection=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, projection=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(projection=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 projection=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(projection=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:
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`)."""
137 return self._mask
139 @property
140 def variance(self) -> Image:
141 """The variance plane (`~lsst.images.Image`)."""
142 return self._variance
144 @property
145 def bbox(self) -> Box:
146 """The bounding box shared by all three image planes
147 (`~lsst.images.Box`).
148 """
149 return self._image.bbox
151 @property
152 def unit(self) -> astropy.units.UnitBase | None:
153 """The units of the image plane (`astropy.units.Unit` | `None`)."""
154 return self._image.unit
156 @property
157 def projection(self) -> Projection[Any] | None:
158 """The projection that maps the pixel grid to the sky
159 (`~lsst.images.Projection` | `None`).
160 """
161 return self._image.projection
163 def __getitem__(self, bbox: Box | EllipsisType) -> MaskedImage:
164 if bbox is ...:
165 return self
166 super().__getitem__(bbox)
167 return self._transfer_metadata(
168 MaskedImage(
169 # Projection and obs_info propagate from the image.
170 self.image[bbox],
171 mask=self.mask[bbox],
172 variance=self.variance[bbox],
173 ),
174 bbox=bbox,
175 )
177 def __setitem__(self, bbox: Box | EllipsisType, value: MaskedImage) -> None:
178 self._image[bbox] = value.image
179 self._mask[bbox] = value.mask
180 self._variance[bbox] = value.variance
182 def __str__(self) -> str:
183 return f"MaskedImage({self.image!s}, {list(self.mask.schema.names)})"
185 def __repr__(self) -> str:
186 return f"MaskedImage({self.image!r}, mask_schema={self.mask.schema!r})"
188 def copy(self) -> MaskedImage:
189 """Deep-copy the masked image and metadata."""
190 return self._transfer_metadata(
191 MaskedImage(image=self._image.copy(), mask=self._mask.copy(), variance=self._variance.copy()),
192 copy=True,
193 )
195 def serialize(self, archive: OutputArchive[Any]) -> MaskedImageSerializationModel:
196 """Serialize the masked image to an output archive.
198 Parameters
199 ----------
200 archive
201 Archive to write to.
202 """
203 return self._serialize_impl(MaskedImageSerializationModel, archive)
205 def _serialize_impl[M: MaskedImageSerializationModel](
206 self, model_type: type[M], archive: OutputArchive[Any]
207 ) -> M:
208 serialized_image = archive.serialize_direct(
209 "image", functools.partial(self.image.serialize, save_projection=False)
210 )
211 serialized_mask = archive.serialize_direct(
212 "mask", functools.partial(self.mask.serialize, save_projection=False)
213 )
214 serialized_variance = archive.serialize_direct(
215 "variance", functools.partial(self.variance.serialize, save_projection=False)
216 )
217 serialized_projection = (
218 archive.serialize_direct("projection", self.projection.serialize)
219 if self.projection is not None
220 else None
221 )
222 # When M is a subclass of MaskedImageSerializationModel, it probably
223 # has fields that aren't being set here. We're intentionally making use
224 # of the fact that model_construct doesn't guard against that so we can
225 # instead set them in a subclass implementation later, after calling
226 # super() to construct an instance of the right type. MyPy is actually
227 # fine with this, but only because it incorrectly thinks model_type is
228 # only ever MaskedImageSerializationModel, not a subclass, and that
229 # makes it incorrectly unhappy about the return type.
230 return model_type.model_construct( # type: ignore[return-value]
231 image=serialized_image,
232 mask=serialized_mask,
233 variance=serialized_variance,
234 projection=serialized_projection,
235 metadata=self.metadata,
236 )
238 @staticmethod
239 def _get_archive_tree_type[P: pydantic.BaseModel](
240 pointer_type: type[P],
241 ) -> type[MaskedImageSerializationModel[P]]:
242 """Return the serialization model type for this object for an archive
243 type that uses the given pointer type.
244 """
245 return MaskedImageSerializationModel[pointer_type] # type: ignore
247 @staticmethod
248 def from_legacy(
249 legacy: LegacyMaskedImage,
250 *,
251 unit: astropy.units.UnitBase | None = None,
252 plane_map: Mapping[str, MaskPlane] | None = None,
253 ) -> MaskedImage:
254 """Convert from an `lsst.afw.image.MaskedImage` instance.
256 Parameters
257 ----------
258 legacy
259 An `lsst.afw.image.MaskedImage` instance that will share image and
260 variance (but not mask) pixel data with the returned object.
261 unit
262 Units of the image.
263 plane_map
264 A mapping from legacy mask plane name to the new plane name and
265 description. If not provided, the right legacy mask plane will be
266 guessed, but this can depend on which mask planes the legacy
267 mask actually has set.
268 """
269 return MaskedImage(
270 image=Image.from_legacy(legacy.getImage(), unit),
271 mask=Mask.from_legacy(legacy.getMask(), plane_map),
272 variance=Image.from_legacy(legacy.getVariance()),
273 )
275 def to_legacy(
276 self, *, copy: bool | None = None, plane_map: Mapping[str, MaskPlane] | None = None
277 ) -> LegacyMaskedImage:
278 """Convert to an `lsst.afw.image.MaskedImage` instance.
280 Parameters
281 ----------
282 copy
283 If `True`, always copy the image and variance pixel data.
284 If `False`, return a view, and raise `TypeError` if the pixel data
285 is read-only (this is not supported by afw). If `None`, only copy
286 if the pixel data is read-only. Mask pixel data is always copied.
287 plane_map
288 A mapping from legacy mask plane name to the new plane name and
289 description.
290 """
291 import lsst.afw.image
293 return lsst.afw.image.MaskedImage(
294 self.image.to_legacy(copy=copy),
295 mask=self.mask.to_legacy(plane_map),
296 variance=self.variance.to_legacy(copy=copy),
297 dtype=self.image.array.dtype,
298 )
300 @staticmethod
301 def read_legacy(
302 uri: ResourcePathExpression,
303 *,
304 preserve_quantization: bool = False,
305 plane_map: Mapping[str, MaskPlane] | None = None,
306 component: Literal["image", "mask", "variance"] | None = None,
307 fits_wcs_frame: Frame | None = None,
308 ) -> Any:
309 """Read a FITS file written by `lsst.afw.image.MaskedImage.writeFits`.
311 Parameters
312 ----------
313 uri
314 URI or file name.
315 preserve_quantization
316 If `True`, ensure that writing the masked image back out again will
317 exactly preserve quantization-compressed pixel values. This causes
318 the image and variance plane arrays to be marked as read-only and
319 stores the original binary table data for those planes in memory.
320 If the `~lsst.images.MaskedImage` is copied, the precompressed
321 pixel values are not transferred to the copy.
322 plane_map
323 A mapping from legacy mask plane name to the new plane name and
324 description. If not provided, the right legacy mask plane will be
325 guessed, but this can depend on which mask planes the legacy
326 mask actually has set.
327 component
328 A component to read instead of the full image.
329 fits_wcs_frame
330 If not `None` and the HDU containing the image plane has a FITS
331 WCS, attach a `~lsst.images.Projection` to the returned masked
332 image by converting that WCS. When ``component`` is one of
333 ``"image"``, ``"mask"``, or ``"variance"``, a FITS WCS from the
334 component HDU is used instead (all three should have the same WCS).
335 """
336 fs, fspath = ResourcePath(uri).to_fsspec()
337 with fs.open(fspath) as stream, astropy.io.fits.open(stream) as hdu_list:
338 return MaskedImage._read_legacy_hdus(
339 hdu_list,
340 uri,
341 preserve_quantization=preserve_quantization,
342 plane_map=plane_map,
343 component=component,
344 fits_wcs_frame=fits_wcs_frame,
345 )
347 @staticmethod
348 def _read_legacy_hdus(
349 hdu_list: astropy.io.fits.HDUList,
350 uri: ResourcePathExpression,
351 *,
352 opaque_metadata: fits.FitsOpaqueMetadata | None = None,
353 preserve_quantization: bool = False,
354 plane_map: Mapping[str, MaskPlane] | None = None,
355 component: Literal["image", "mask", "variance"] | None,
356 fits_wcs_frame: Frame | None = None,
357 ) -> Any:
358 if opaque_metadata is None:
359 opaque_metadata = fits.FitsOpaqueMetadata()
360 opaque_metadata.extract_legacy_primary_header(hdu_list[0].header)
361 image_bintable_hdu: astropy.io.fits.BinTableHDU | None = None
362 variance_bintable_hdu: astropy.io.fits.BinTableHDU | None = None
363 result: Any
364 with ExitStack() as exit_stack:
365 if preserve_quantization:
366 fs, fspath = ResourcePath(uri).to_fsspec()
367 bintable_stream = exit_stack.enter_context(fs.open(fspath))
368 bintable_hdu_list = exit_stack.enter_context(
369 astropy.io.fits.open(bintable_stream, disable_image_compression=True)
370 )
371 image_bintable_hdu = bintable_hdu_list[1]
372 variance_bintable_hdu = bintable_hdu_list[3]
373 if component is None or component == "image":
374 image = Image._read_legacy_hdu(
375 hdu_list[1],
376 opaque_metadata,
377 preserve_bintable=image_bintable_hdu,
378 fits_wcs_frame=fits_wcs_frame,
379 )
380 if component == "image":
381 result = image
382 if component is None or component == "mask":
383 mask = Mask._read_legacy_hdu(
384 hdu_list[2],
385 opaque_metadata,
386 plane_map=plane_map,
387 fits_wcs_frame=fits_wcs_frame if component is not None else None,
388 )
389 if component == "mask":
390 result = mask
391 if component is None or component == "variance":
392 variance = Image._read_legacy_hdu(
393 hdu_list[3],
394 opaque_metadata,
395 preserve_bintable=variance_bintable_hdu,
396 fits_wcs_frame=fits_wcs_frame if component is not None else None,
397 )
398 if component == "variance":
399 result = variance
400 if component is None:
401 result = MaskedImage(image, mask=mask, variance=variance)
402 result._opaque_metadata = opaque_metadata
403 return result
406class MaskedImageSerializationModel[P: pydantic.BaseModel](ArchiveTree):
407 """A Pydantic model used to represent a serialized `MaskedImage`."""
409 SCHEMA_NAME: ClassVar[str] = "masked_image"
410 SCHEMA_VERSION: ClassVar[str] = "1.0.0"
411 MIN_READ_VERSION: ClassVar[int] = 1
412 PUBLIC_TYPE: ClassVar[type] = MaskedImage
414 image: ImageSerializationModel[P] = pydantic.Field(description="The main data image.")
415 mask: MaskSerializationModel[P] = pydantic.Field(
416 description="Bitmask that annotates the main image's pixels."
417 )
418 variance: ImageSerializationModel[P] = pydantic.Field(
419 description="Per-pixel variance estimates for the main image."
420 )
421 projection: ProjectionSerializationModel[P] | None = pydantic.Field(
422 default=None,
423 exclude_if=is_none,
424 description="Projection that maps the pixel grid to the sky.",
425 )
427 @property
428 def bbox(self) -> Box:
429 """The bounding box of the image."""
430 return self.image.bbox
432 def deserialize(
433 self, archive: InputArchive[Any], *, bbox: Box | None = None, **kwargs: Any
434 ) -> MaskedImage:
435 """Deserialize an image from an input archive.
437 Parameters
438 ----------
439 archive
440 Archive to read from.
441 bbox
442 Bounding box of a subimage to read instead.
443 **kwargs
444 Unsupported keyword arguments are accepted only to provide better
445 error messages (raising `serialization.InvalidParameterError`).
446 """
447 if kwargs:
448 raise InvalidParameterError(f"Unrecognized parameters for MaskedImage: {set(kwargs.keys())}.")
449 image = self.image.deserialize(archive, bbox=bbox)
450 mask = self.mask.deserialize(archive, bbox=bbox)
451 variance = self.variance.deserialize(archive, bbox=bbox)
452 projection = self.projection.deserialize(archive) if self.projection is not None else None
453 return MaskedImage(image, mask=mask, variance=variance, projection=projection)._finish_deserialize(
454 self
455 )
457 def deserialize_component(self, component: str, archive: InputArchive[Any], **kwargs: Any) -> Any:
458 if component == "bbox" and kwargs:
459 raise InvalidParameterError(
460 f"Unrecognized parameters for MaskedImage.bbox: {set(kwargs.keys())}."
461 )
462 return super().deserialize_component(component, archive, **kwargs)