Coverage for python / lsst / images / _masked_image.py: 34%
182 statements
« prev ^ index » next coverage.py v7.13.4, created at 2026-03-02 07:17 -0800
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 typing import Any, Literal, overload
21import astropy.io.fits
22import astropy.units
23import astropy.wcs
24import numpy as np
25import pydantic
26from astro_metadata_translator import ObservationInfo
28from lsst.resources import ResourcePath, ResourcePathExpression
30from . import fits
31from ._geom import Box
32from ._image import Image, ImageSerializationModel
33from ._mask import Mask, MaskPlane, MaskSchema, MaskSerializationModel
34from ._transforms import Frame, Projection, ProjectionAstropyView, ProjectionSerializationModel
35from .serialization import (
36 ArchiveTree,
37 InputArchive,
38 OpaqueArchiveMetadata,
39 OutputArchive,
40)
41from .utils import is_none
44class MaskedImage:
45 """A multi-plane image with data (image), mask, and variance planes.
47 Parameters
48 ----------
49 image
50 The main image plane. If this has a `Projection`, it will be used
51 for all planes unless a ``projection`` is passed separately.
52 mask
53 A bitmask image that annotates the main image plane. Must have the
54 same bounding box as ``image`` if provided. Any attached projection
55 is replaced (possibly by `None`).
56 variance
57 The per-pixel uncertainty of the main image as an image of variance
58 values. Must have the same bounding box as ``image`` if provided, and
59 its units must be the square of ``image.unit`` or `None`.
60 Values default to ``1.0``. Any attached projection is replaced
61 (possibly by `None`).
62 mask_schema
63 Schema for the mask plane. Must be provided if and only if ``mask`` is
64 not provided.
65 projection
66 Projection that maps the pixel grid to the sky.
67 obs_info
68 General information about the associated observation in standardized
69 form.
70 """
72 def __init__(
73 self,
74 image: Image,
75 *,
76 mask: Mask | None = None,
77 variance: Image | None = None,
78 mask_schema: MaskSchema | None = None,
79 projection: Projection | None = None,
80 obs_info: ObservationInfo | None = None,
81 ):
82 if projection is None:
83 projection = image.projection
84 else:
85 image = image.view(projection=projection)
86 if obs_info is None:
87 obs_info = image.obs_info
88 else:
89 image = image.view(obs_info=obs_info)
90 if mask is None:
91 if mask_schema is None:
92 raise TypeError("'mask_schema' must be provided if 'mask' is not.")
93 mask = Mask(schema=mask_schema, bbox=image.bbox, projection=projection, obs_info=obs_info)
94 elif mask_schema is not None:
95 raise TypeError("'mask_schema' may not be provided if 'mask' is.")
96 else:
97 if image.bbox != mask.bbox:
98 raise ValueError(f"Image ({image.bbox}) and mask ({mask.bbox}) bboxes do not agree.")
99 mask = mask.view(projection=projection)
100 if variance is None:
101 variance = Image(
102 1.0,
103 dtype=np.float32,
104 bbox=image.bbox,
105 unit=None if image.unit is None else image.unit**2,
106 projection=projection,
107 obs_info=obs_info,
108 )
109 else:
110 if image.bbox != variance.bbox:
111 raise ValueError(f"Image ({image.bbox}) and variance ({variance.bbox}) bboxes do not agree.")
112 variance = variance.view(projection=projection, obs_info=obs_info)
113 if image.unit is None:
114 if variance.unit is not None:
115 raise ValueError(f"Image has no units but variance does ({variance.unit}).")
116 elif variance.unit is None:
117 variance = variance.view(unit=image.unit**2)
118 elif variance.unit != image.unit**2:
119 raise ValueError(
120 f"Variance unit ({variance.unit}) should be the square of the image unit ({image.unit})."
121 )
122 self._image = image
123 self._mask = mask
124 self._variance = variance
125 self._opaque_metadata: OpaqueArchiveMetadata | None = None
127 @property
128 def image(self) -> Image:
129 """The main image plane (`Image`)."""
130 return self._image
132 @property
133 def mask(self) -> Mask:
134 """The mask plane (`Mask`)."""
135 return self._mask
137 @property
138 def variance(self) -> Image:
139 """The variance plane (`Image`)."""
140 return self._variance
142 @property
143 def bbox(self) -> Box:
144 """The bounding box shared by all three image planes (`Box`)."""
145 return self._image.bbox
147 @property
148 def unit(self) -> astropy.units.UnitBase | None:
149 """The units of the image plane (`astropy.units.Unit` | `None`)."""
150 return self._image.unit
152 @property
153 def projection(self) -> Projection[Any] | None:
154 """The projection that maps the pixel grid to the sky
155 (`Projection` | `None`).
156 """
157 return self._image.projection
159 @property
160 def obs_info(self) -> ObservationInfo | None:
161 """General information about the associated observation in standard
162 form. (`~astro_metadata_translator.ObservationInfo` | `None`).
163 """
164 return self._image.obs_info
166 @property
167 def astropy_wcs(self) -> ProjectionAstropyView | None:
168 """An Astropy WCS for the pixel arrays
169 (`ProjectionAstropyView` | `None`).
171 Notes
172 -----
173 As expected for Astropy WCS objects, this defines pixel coordinates
174 such that the first row and column in the arrays are ``(0, 0)``, not
175 ``bbox.start``, as is the case for `projection`.
177 This object satisfies the `astropy.wcs.wcsapi.BaseHighLevelWCS` and
178 `astropy.wcs.wcsapi.BaseLowLevelWCS` interfaces, but it is not an
179 `astropy.wcs.WCS` (use `fits_wcs` for that).
180 """
181 return self.projection.as_astropy(self.bbox) if self.projection is not None else None
183 @property
184 def fits_wcs(self) -> astropy.wcs.WCS | None:
185 """An Astropy FITS WCS for this mask's pixel array
186 (`astropy.wcs.WCS` | `None`).
188 Notes
189 -----
190 As expected for Astropy WCS objects, this defines pixel coordinates
191 such that the first row and column in the arrays are ``(0, 0)``, not
192 ``bbox.start``, as is the case for `projection`.
194 This may be an approximation or absent if `projection` is not
195 naturally representable as a FITS WCS.
196 """
197 return self.projection.as_fits_wcs(self.bbox) if self.projection is not None else None
199 def __getitem__(self, bbox: Box) -> MaskedImage:
200 result = MaskedImage(
201 # Projection and obs_info propagate from the image.
202 self.image[bbox],
203 mask=self.mask[bbox],
204 variance=self.variance[bbox],
205 )
206 if self._opaque_metadata is not None:
207 result._opaque_metadata = self._opaque_metadata.subset(bbox)
208 return result
210 def __str__(self) -> str:
211 return f"MaskedImage({self.image!s}, {list(self.mask.schema.names)})"
213 def __repr__(self) -> str:
214 return f"MaskedImage({self.image!r}, mask_schema={self.mask.schema!r})"
216 def copy(self) -> MaskedImage:
217 """Deep-copy the masked image."""
218 result = MaskedImage(image=self._image.copy(), mask=self._mask.copy(), variance=self._variance.copy())
219 if self._opaque_metadata is not None:
220 result._opaque_metadata = self._opaque_metadata.copy()
221 return result
223 def serialize(self, archive: OutputArchive[Any]) -> MaskedImageSerializationModel:
224 """Serialize the masked image to an output archive.
226 Parameters
227 ----------
228 archive
229 Archive to write to.
230 """
231 serialized_image = archive.serialize_direct(
232 "image", functools.partial(self.image.serialize, save_projection=False)
233 )
234 serialized_mask = archive.serialize_direct(
235 "mask", functools.partial(self.mask.serialize, save_projection=False)
236 )
237 serialized_variance = archive.serialize_direct(
238 "variance", functools.partial(self.variance.serialize, save_projection=False)
239 )
240 serialized_projection = (
241 archive.serialize_direct("projection", self.projection.serialize)
242 if self.projection is not None
243 else None
244 )
245 return MaskedImageSerializationModel(
246 image=serialized_image,
247 mask=serialized_mask,
248 variance=serialized_variance,
249 projection=serialized_projection,
250 )
252 @staticmethod
253 def deserialize(
254 model: MaskedImageSerializationModel[Any], archive: InputArchive[Any], *, bbox: Box | None = None
255 ) -> MaskedImage:
256 """Deserialize an image from an input archive.
258 Parameters
259 ----------
260 model
261 A Pydantic model representation of the image, holding references
262 to data stored in the archive.
263 archive
264 Archive to read from.
265 bbox
266 Bounding box of a subimage to read instead.
267 """
268 image = Image.deserialize(model.image, archive, bbox=bbox)
269 mask = Mask.deserialize(model.mask, archive, bbox=bbox)
270 variance = Image.deserialize(model.variance, archive, bbox=bbox)
271 projection = (
272 Projection.deserialize(model.projection, archive) if model.projection is not None else None
273 )
274 return MaskedImage(image, mask=mask, variance=variance, projection=projection)
276 @staticmethod
277 def _get_archive_tree_type[P: pydantic.BaseModel](
278 pointer_type: type[P],
279 ) -> type[MaskedImageSerializationModel[P]]:
280 """Return the serialization model type for this object for an archive
281 type that uses the given pointer type.
282 """
283 return MaskedImageSerializationModel[pointer_type] # type: ignore
285 def write_fits(
286 self,
287 filename: str,
288 *,
289 image_compression: fits.FitsCompressionOptions | None = fits.FitsCompressionOptions.DEFAULT,
290 mask_compression: fits.FitsCompressionOptions | None = fits.FitsCompressionOptions.DEFAULT,
291 variance_compression: fits.FitsCompressionOptions | None = fits.FitsCompressionOptions.DEFAULT,
292 compression_seed: int | None = None,
293 ) -> None:
294 """Write the image to a FITS file.
296 Parameters
297 ----------
298 filename
299 Name of the file to write to. Must be a local file.
300 image_compression
301 Compression options for the `image` plane.
302 mask_compression
303 Compression options for the `mask` plane.
304 variance_compression
305 Compression options for the `variance` plane.
306 compression_seed
307 A FITS tile compression seed to use whenever the configured
308 compression seed is `None` or (for backwards compatibility) ``0``.
309 This value is then incremented every time it is used.
310 """
311 compression_options = {}
312 if image_compression is not fits.FitsCompressionOptions.DEFAULT:
313 compression_options["image"] = image_compression
314 if mask_compression is not fits.FitsCompressionOptions.DEFAULT:
315 compression_options["mask"] = mask_compression
316 if variance_compression is not fits.FitsCompressionOptions.DEFAULT:
317 compression_options["variance"] = variance_compression
318 fits.write(self, filename, compression_options=compression_options, compression_seed=compression_seed)
320 @classmethod
321 def read_fits(cls, url: ResourcePathExpression, *, bbox: Box | None = None) -> MaskedImage:
322 """Read an image from a FITS file.
324 Parameters
325 ----------
326 url
327 URL of the file to read; may be any type supported by
328 `lsst.resources.ResourcePath`.
329 bbox
330 Bounding box of a subimage to read instead.
331 """
332 return fits.read(cls, url, bbox=bbox)
334 @staticmethod
335 def from_legacy(
336 legacy: Any,
337 *,
338 unit: astropy.units.Unit | None = None,
339 plane_map: Mapping[str, MaskPlane] | None = None,
340 ) -> MaskedImage:
341 """Convert from an `lsst.afw.image.MaskedImage` instance.
343 Parameters
344 ----------
345 legacy
346 An `lsst.afw.image.MaskedImage` instance that will share image and
347 variance (but not mask) pixel data with the returned object.
348 unit
349 Units of the image.
350 plane_map
351 A mapping from legacy mask plane name to the new plane name and
352 description.
353 """
354 return MaskedImage(
355 image=Image.from_legacy(legacy.getImage(), unit),
356 mask=Mask.from_legacy(legacy.getMask(), plane_map),
357 variance=Image.from_legacy(legacy.getVariance()),
358 )
360 def to_legacy(self, *, copy: bool | None = None, plane_map: Mapping[str, MaskPlane] | None = None) -> Any:
361 """Convert to an `lsst.afw.image.MaskedImage` instance.
363 Parameters
364 ----------
365 copy
366 If `True`, always copy the image and variance pixel data.
367 If `False`, return a view, and raise `TypeError` if the pixel data
368 is read-only (this is not supported by afw). If `None`, onyl if
369 the pixel data is read-only. Mask pixel data is always copied.
370 plane_map
371 A mapping from legacy mask plane name to the new plane name and
372 description.
373 """
374 import lsst.afw.image
376 return lsst.afw.image.MaskedImage(
377 self.image.to_legacy(copy=copy),
378 mask=self.mask.to_legacy(plane_map),
379 variance=self.variance.to_legacy(copy=copy),
380 dtype=self.image.array.dtype,
381 )
383 @overload
384 @staticmethod
385 def read_legacy( 385 ↛ exitline 385 didn't return from function 'read_legacy' because
386 uri: ResourcePathExpression,
387 *,
388 preserve_quantization: bool = False,
389 component: Literal["image"],
390 fits_wcs_frame: Frame | None = None,
391 ) -> Image: ...
393 @overload
394 @staticmethod
395 def read_legacy( 395 ↛ exitline 395 didn't return from function 'read_legacy' because
396 uri: ResourcePathExpression,
397 *,
398 plane_map: Mapping[str, MaskPlane] | None = None,
399 component: Literal["mask"],
400 fits_wcs_frame: Frame | None = None,
401 ) -> Mask: ...
403 @overload
404 @staticmethod
405 def read_legacy( 405 ↛ exitline 405 didn't return from function 'read_legacy' because
406 uri: ResourcePathExpression,
407 *,
408 preserve_quantization: bool = False,
409 component: Literal["variance"],
410 fits_wcs_frame: Frame | None = None,
411 ) -> Image: ...
413 @overload
414 @staticmethod
415 def read_legacy( 415 ↛ exitline 415 didn't return from function 'read_legacy' because
416 uri: ResourcePathExpression,
417 *,
418 preserve_quantization: bool = False,
419 plane_map: Mapping[str, MaskPlane] | None = None,
420 component: None = None,
421 fits_wcs_frame: Frame | None = None,
422 ) -> MaskedImage: ...
424 @staticmethod
425 def read_legacy(
426 uri: ResourcePathExpression,
427 *,
428 preserve_quantization: bool = False,
429 plane_map: Mapping[str, MaskPlane] | None = None,
430 component: Literal["image", "mask", "variance"] | None = None,
431 fits_wcs_frame: Frame | None = None,
432 ) -> Any:
433 """Read a FITS file written by `lsst.afw.image.MaskedImage.writeFits`.
435 Parameters
436 ----------
437 uri
438 URI or file name.
439 preserve_quantization
440 If `True`, ensure that writing the masked image back out again will
441 exactly preserve quantization-compressed pixel values. This causes
442 the image and variance plane arrays to be marked as read-only and
443 stores the original binary table data for those planes in memory.
444 If the `MaskedImage` is copied, the precompressed pixel values are
445 not transferred to the copy.
446 plane_map
447 A mapping from legacy mask plane name to the new plane name and
448 description.
449 component
450 A component to read instead of the full image.
451 fits_wcs_frame
452 If not `None` and the HDU containing the image plane has a FITS
453 WCS, attach a `Projection` to the returned masked image by
454 converting that WCS. When ``component`` is one of ``"image"``,
455 ``"mask"``, or ``"variance"``, a FITS WCS from the component HDU
456 is used instead (all three should have the same WCS).
457 """
458 fs, fspath = ResourcePath(uri).to_fsspec()
459 with fs.open(fspath) as stream, astropy.io.fits.open(stream) as hdu_list:
460 return MaskedImage._read_legacy_hdus(
461 hdu_list,
462 uri,
463 preserve_quantization=preserve_quantization,
464 plane_map=plane_map,
465 component=component,
466 fits_wcs_frame=fits_wcs_frame,
467 )
469 @staticmethod
470 def _read_legacy_hdus(
471 hdu_list: astropy.io.fits.HDUList,
472 uri: ResourcePathExpression,
473 *,
474 preserve_quantization: bool = False,
475 plane_map: Mapping[str, MaskPlane] | None = None,
476 component: Literal["image", "mask", "variance"] | None,
477 fits_wcs_frame: Frame | None = None,
478 ) -> Any:
479 opaque_metadata = fits.FitsOpaqueMetadata()
480 opaque_metadata.extract_legacy_primary_header(hdu_list[0].header)
481 image_bintable_hdu: astropy.io.fits.BinTableHDU | None = None
482 variance_bintable_hdu: astropy.io.fits.BinTableHDU | None = None
483 result: Any
484 with ExitStack() as exit_stack:
485 if preserve_quantization:
486 fs, fspath = ResourcePath(uri).to_fsspec()
487 bintable_stream = exit_stack.enter_context(fs.open(fspath))
488 bintable_hdu_list = exit_stack.enter_context(
489 astropy.io.fits.open(bintable_stream, disable_image_compression=True)
490 )
491 image_bintable_hdu = bintable_hdu_list[1]
492 variance_bintable_hdu = bintable_hdu_list[3]
493 if component is None or component == "image":
494 image = Image._read_legacy_hdu(
495 hdu_list[1],
496 opaque_metadata,
497 preserve_bintable=image_bintable_hdu,
498 fits_wcs_frame=fits_wcs_frame,
499 )
500 if component == "image":
501 result = image
502 if component is None or component == "mask":
503 mask = Mask._read_legacy_hdu(
504 hdu_list[2],
505 opaque_metadata,
506 plane_map=plane_map,
507 fits_wcs_frame=fits_wcs_frame if component is not None else None,
508 )
509 if component == "mask":
510 result = mask
511 if component is None or component == "variance":
512 variance = Image._read_legacy_hdu(
513 hdu_list[3],
514 opaque_metadata,
515 preserve_bintable=variance_bintable_hdu,
516 fits_wcs_frame=fits_wcs_frame if component is not None else None,
517 )
518 if component == "variance":
519 result = variance
520 if component is None:
521 result = MaskedImage(image, mask=mask, variance=variance)
522 result._opaque_metadata = opaque_metadata
523 return result
526class MaskedImageSerializationModel[P: pydantic.BaseModel](ArchiveTree):
527 """A Pydantic model used to represent a serialized `MaskedImage`."""
529 image: ImageSerializationModel[P] = pydantic.Field(description="The main data image.")
530 mask: MaskSerializationModel[P] = pydantic.Field(
531 description="Bitmask that annotates the main image's pixels."
532 )
533 variance: ImageSerializationModel[P] = pydantic.Field(
534 description="Per-pixel variance estimates for the main image."
535 )
536 projection: ProjectionSerializationModel[P] | None = pydantic.Field(
537 default=None,
538 exclude_if=is_none,
539 description="Projection that maps the pixel grid to the sky.",
540 )
542 @property
543 def bbox(self) -> Box:
544 """The bounding box of the image."""
545 return self.image.bbox