Coverage for python / lsst / images / _masked_image.py: 33%
185 statements
« prev ^ index » next coverage.py v7.13.4, created at 2026-03-05 02:32 -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 types import EllipsisType
20from typing import Any, Literal, overload
22import astropy.io.fits
23import astropy.units
24import astropy.wcs
25import numpy as np
26import pydantic
27from astro_metadata_translator import ObservationInfo
29from lsst.resources import ResourcePath, ResourcePathExpression
31from . import fits
32from ._generalized_image import GeneralizedImage
33from ._geom import Box
34from ._image import Image, ImageSerializationModel
35from ._mask import Mask, MaskPlane, MaskSchema, MaskSerializationModel
36from ._transforms import Frame, Projection, ProjectionSerializationModel
37from .serialization import (
38 ArchiveTree,
39 InputArchive,
40 OpaqueArchiveMetadata,
41 OutputArchive,
42)
43from .utils import is_none
46class MaskedImage(GeneralizedImage):
47 """A multi-plane image with data (image), mask, and variance planes.
49 Parameters
50 ----------
51 image
52 The main image plane. If this has a `Projection`, it will be used
53 for all planes unless a ``projection`` is passed separately.
54 mask
55 A bitmask image that annotates the main image plane. Must have the
56 same bounding box as ``image`` if provided. Any attached projection
57 is replaced (possibly by `None`).
58 variance
59 The per-pixel uncertainty of the main image as an image of variance
60 values. Must have the same bounding box as ``image`` if provided, and
61 its units must be the square of ``image.unit`` or `None`.
62 Values default to ``1.0``. Any attached projection is replaced
63 (possibly by `None`).
64 mask_schema
65 Schema for the mask plane. Must be provided if and only if ``mask`` is
66 not provided.
67 projection
68 Projection that maps the pixel grid to the sky.
69 obs_info
70 General information about the associated observation in standardized
71 form.
72 """
74 def __init__(
75 self,
76 image: Image,
77 *,
78 mask: Mask | None = None,
79 variance: Image | None = None,
80 mask_schema: MaskSchema | None = None,
81 projection: Projection | None = None,
82 obs_info: ObservationInfo | None = None,
83 ):
84 if projection is None:
85 projection = image.projection
86 else:
87 image = image.view(projection=projection)
88 if obs_info is None:
89 obs_info = image.obs_info
90 else:
91 image = image.view(obs_info=obs_info)
92 if mask is None:
93 if mask_schema is None:
94 raise TypeError("'mask_schema' must be provided if 'mask' is not.")
95 mask = Mask(schema=mask_schema, bbox=image.bbox, projection=projection, obs_info=obs_info)
96 elif mask_schema is not None:
97 raise TypeError("'mask_schema' may not be provided if 'mask' is.")
98 else:
99 if image.bbox != mask.bbox:
100 raise ValueError(f"Image ({image.bbox}) and mask ({mask.bbox}) bboxes do not agree.")
101 mask = mask.view(projection=projection)
102 if variance is None:
103 variance = Image(
104 1.0,
105 dtype=np.float32,
106 bbox=image.bbox,
107 unit=None if image.unit is None else image.unit**2,
108 projection=projection,
109 obs_info=obs_info,
110 )
111 else:
112 if image.bbox != variance.bbox:
113 raise ValueError(f"Image ({image.bbox}) and variance ({variance.bbox}) bboxes do not agree.")
114 variance = variance.view(projection=projection, obs_info=obs_info)
115 if image.unit is None:
116 if variance.unit is not None:
117 raise ValueError(f"Image has no units but variance does ({variance.unit}).")
118 elif variance.unit is None:
119 variance = variance.view(unit=image.unit**2)
120 elif variance.unit != image.unit**2:
121 raise ValueError(
122 f"Variance unit ({variance.unit}) should be the square of the image unit ({image.unit})."
123 )
124 self._image = image
125 self._mask = mask
126 self._variance = variance
127 self._opaque_metadata: OpaqueArchiveMetadata | None = None
129 @property
130 def image(self) -> Image:
131 """The main image plane (`Image`)."""
132 return self._image
134 @property
135 def mask(self) -> Mask:
136 """The mask plane (`Mask`)."""
137 return self._mask
139 @property
140 def variance(self) -> Image:
141 """The variance plane (`Image`)."""
142 return self._variance
144 @property
145 def bbox(self) -> Box:
146 """The bounding box shared by all three image planes (`Box`)."""
147 return self._image.bbox
149 @property
150 def unit(self) -> astropy.units.UnitBase | None:
151 """The units of the image plane (`astropy.units.Unit` | `None`)."""
152 return self._image.unit
154 @property
155 def projection(self) -> Projection[Any] | None:
156 """The projection that maps the pixel grid to the sky
157 (`Projection` | `None`).
158 """
159 return self._image.projection
161 @property
162 def obs_info(self) -> ObservationInfo | None:
163 """General information about the associated observation in standard
164 form. (`~astro_metadata_translator.ObservationInfo` | `None`).
165 """
166 return self._image.obs_info
168 def __getitem__(self, bbox: Box | EllipsisType) -> MaskedImage:
169 super().__getitem__(bbox)
170 if bbox is ...:
171 return self
172 result = MaskedImage(
173 # Projection and obs_info propagate from the image.
174 self.image[bbox],
175 mask=self.mask[bbox],
176 variance=self.variance[bbox],
177 )
178 if self._opaque_metadata is not None:
179 result._opaque_metadata = self._opaque_metadata.subset(bbox)
180 return result
182 def __setitem__(self, bbox: Box | EllipsisType, value: MaskedImage) -> None:
183 self._image[bbox] = value.image
184 self._mask[bbox] = value.mask
185 self._variance[bbox] = value.variance
187 def __str__(self) -> str:
188 return f"MaskedImage({self.image!s}, {list(self.mask.schema.names)})"
190 def __repr__(self) -> str:
191 return f"MaskedImage({self.image!r}, mask_schema={self.mask.schema!r})"
193 def copy(self) -> MaskedImage:
194 """Deep-copy the masked image."""
195 result = MaskedImage(image=self._image.copy(), mask=self._mask.copy(), variance=self._variance.copy())
196 if self._opaque_metadata is not None:
197 result._opaque_metadata = self._opaque_metadata.copy()
198 return result
200 def serialize(self, archive: OutputArchive[Any]) -> MaskedImageSerializationModel:
201 """Serialize the masked image to an output archive.
203 Parameters
204 ----------
205 archive
206 Archive to write to.
207 """
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 return MaskedImageSerializationModel(
223 image=serialized_image,
224 mask=serialized_mask,
225 variance=serialized_variance,
226 projection=serialized_projection,
227 )
229 @staticmethod
230 def deserialize(
231 model: MaskedImageSerializationModel[Any], archive: InputArchive[Any], *, bbox: Box | None = None
232 ) -> MaskedImage:
233 """Deserialize an image from an input archive.
235 Parameters
236 ----------
237 model
238 A Pydantic model representation of the image, holding references
239 to data stored in the archive.
240 archive
241 Archive to read from.
242 bbox
243 Bounding box of a subimage to read instead.
244 """
245 image = Image.deserialize(model.image, archive, bbox=bbox)
246 mask = Mask.deserialize(model.mask, archive, bbox=bbox)
247 variance = Image.deserialize(model.variance, archive, bbox=bbox)
248 projection = (
249 Projection.deserialize(model.projection, archive) if model.projection is not None else None
250 )
251 return MaskedImage(image, mask=mask, variance=variance, projection=projection)
253 @staticmethod
254 def _get_archive_tree_type[P: pydantic.BaseModel](
255 pointer_type: type[P],
256 ) -> type[MaskedImageSerializationModel[P]]:
257 """Return the serialization model type for this object for an archive
258 type that uses the given pointer type.
259 """
260 return MaskedImageSerializationModel[pointer_type] # type: ignore
262 def write_fits(
263 self,
264 filename: str,
265 *,
266 image_compression: fits.FitsCompressionOptions | None = fits.FitsCompressionOptions.DEFAULT,
267 mask_compression: fits.FitsCompressionOptions | None = fits.FitsCompressionOptions.DEFAULT,
268 variance_compression: fits.FitsCompressionOptions | None = fits.FitsCompressionOptions.DEFAULT,
269 compression_seed: int | None = None,
270 ) -> None:
271 """Write the image to a FITS file.
273 Parameters
274 ----------
275 filename
276 Name of the file to write to. Must be a local file.
277 image_compression
278 Compression options for the `image` plane.
279 mask_compression
280 Compression options for the `mask` plane.
281 variance_compression
282 Compression options for the `variance` plane.
283 compression_seed
284 A FITS tile compression seed to use whenever the configured
285 compression seed is `None` or (for backwards compatibility) ``0``.
286 This value is then incremented every time it is used.
287 """
288 compression_options = {}
289 if image_compression is not fits.FitsCompressionOptions.DEFAULT:
290 compression_options["image"] = image_compression
291 if mask_compression is not fits.FitsCompressionOptions.DEFAULT:
292 compression_options["mask"] = mask_compression
293 if variance_compression is not fits.FitsCompressionOptions.DEFAULT:
294 compression_options["variance"] = variance_compression
295 fits.write(self, filename, compression_options=compression_options, compression_seed=compression_seed)
297 @classmethod
298 def read_fits(cls, url: ResourcePathExpression, *, bbox: Box | None = None) -> MaskedImage:
299 """Read an image from a FITS file.
301 Parameters
302 ----------
303 url
304 URL of the file to read; may be any type supported by
305 `lsst.resources.ResourcePath`.
306 bbox
307 Bounding box of a subimage to read instead.
308 """
309 return fits.read(cls, url, bbox=bbox)
311 @staticmethod
312 def from_legacy(
313 legacy: Any,
314 *,
315 unit: astropy.units.Unit | None = None,
316 plane_map: Mapping[str, MaskPlane] | None = None,
317 ) -> MaskedImage:
318 """Convert from an `lsst.afw.image.MaskedImage` instance.
320 Parameters
321 ----------
322 legacy
323 An `lsst.afw.image.MaskedImage` instance that will share image and
324 variance (but not mask) pixel data with the returned object.
325 unit
326 Units of the image.
327 plane_map
328 A mapping from legacy mask plane name to the new plane name and
329 description.
330 """
331 return MaskedImage(
332 image=Image.from_legacy(legacy.getImage(), unit),
333 mask=Mask.from_legacy(legacy.getMask(), plane_map),
334 variance=Image.from_legacy(legacy.getVariance()),
335 )
337 def to_legacy(self, *, copy: bool | None = None, plane_map: Mapping[str, MaskPlane] | None = None) -> Any:
338 """Convert to an `lsst.afw.image.MaskedImage` instance.
340 Parameters
341 ----------
342 copy
343 If `True`, always copy the image and variance pixel data.
344 If `False`, return a view, and raise `TypeError` if the pixel data
345 is read-only (this is not supported by afw). If `None`, onyl if
346 the pixel data is read-only. Mask pixel data is always copied.
347 plane_map
348 A mapping from legacy mask plane name to the new plane name and
349 description.
350 """
351 import lsst.afw.image
353 return lsst.afw.image.MaskedImage(
354 self.image.to_legacy(copy=copy),
355 mask=self.mask.to_legacy(plane_map),
356 variance=self.variance.to_legacy(copy=copy),
357 dtype=self.image.array.dtype,
358 )
360 @overload
361 @staticmethod
362 def read_legacy( 362 ↛ exitline 362 didn't return from function 'read_legacy' because
363 uri: ResourcePathExpression,
364 *,
365 preserve_quantization: bool = False,
366 component: Literal["image"],
367 fits_wcs_frame: Frame | None = None,
368 ) -> Image: ...
370 @overload
371 @staticmethod
372 def read_legacy( 372 ↛ exitline 372 didn't return from function 'read_legacy' because
373 uri: ResourcePathExpression,
374 *,
375 plane_map: Mapping[str, MaskPlane] | None = None,
376 component: Literal["mask"],
377 fits_wcs_frame: Frame | None = None,
378 ) -> Mask: ...
380 @overload
381 @staticmethod
382 def read_legacy( 382 ↛ exitline 382 didn't return from function 'read_legacy' because
383 uri: ResourcePathExpression,
384 *,
385 preserve_quantization: bool = False,
386 component: Literal["variance"],
387 fits_wcs_frame: Frame | None = None,
388 ) -> Image: ...
390 @overload
391 @staticmethod
392 def read_legacy( 392 ↛ exitline 392 didn't return from function 'read_legacy' because
393 uri: ResourcePathExpression,
394 *,
395 preserve_quantization: bool = False,
396 plane_map: Mapping[str, MaskPlane] | None = None,
397 component: None = None,
398 fits_wcs_frame: Frame | None = None,
399 ) -> MaskedImage: ...
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 `MaskedImage` is copied, the precompressed pixel values are
422 not transferred to the copy.
423 plane_map
424 A mapping from legacy mask plane name to the new plane name and
425 description.
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 `Projection` to the returned masked image by
431 converting that WCS. When ``component`` is one of ``"image"``,
432 ``"mask"``, or ``"variance"``, a FITS WCS from the component HDU
433 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 preserve_quantization: bool = False,
452 plane_map: Mapping[str, MaskPlane] | None = None,
453 component: Literal["image", "mask", "variance"] | None,
454 fits_wcs_frame: Frame | None = None,
455 ) -> Any:
456 opaque_metadata = fits.FitsOpaqueMetadata()
457 opaque_metadata.extract_legacy_primary_header(hdu_list[0].header)
458 image_bintable_hdu: astropy.io.fits.BinTableHDU | None = None
459 variance_bintable_hdu: astropy.io.fits.BinTableHDU | None = None
460 result: Any
461 with ExitStack() as exit_stack:
462 if preserve_quantization:
463 fs, fspath = ResourcePath(uri).to_fsspec()
464 bintable_stream = exit_stack.enter_context(fs.open(fspath))
465 bintable_hdu_list = exit_stack.enter_context(
466 astropy.io.fits.open(bintable_stream, disable_image_compression=True)
467 )
468 image_bintable_hdu = bintable_hdu_list[1]
469 variance_bintable_hdu = bintable_hdu_list[3]
470 if component is None or component == "image":
471 image = Image._read_legacy_hdu(
472 hdu_list[1],
473 opaque_metadata,
474 preserve_bintable=image_bintable_hdu,
475 fits_wcs_frame=fits_wcs_frame,
476 )
477 if component == "image":
478 result = image
479 if component is None or component == "mask":
480 mask = Mask._read_legacy_hdu(
481 hdu_list[2],
482 opaque_metadata,
483 plane_map=plane_map,
484 fits_wcs_frame=fits_wcs_frame if component is not None else None,
485 )
486 if component == "mask":
487 result = mask
488 if component is None or component == "variance":
489 variance = Image._read_legacy_hdu(
490 hdu_list[3],
491 opaque_metadata,
492 preserve_bintable=variance_bintable_hdu,
493 fits_wcs_frame=fits_wcs_frame if component is not None else None,
494 )
495 if component == "variance":
496 result = variance
497 if component is None:
498 result = MaskedImage(image, mask=mask, variance=variance)
499 result._opaque_metadata = opaque_metadata
500 return result
503class MaskedImageSerializationModel[P: pydantic.BaseModel](ArchiveTree):
504 """A Pydantic model used to represent a serialized `MaskedImage`."""
506 image: ImageSerializationModel[P] = pydantic.Field(description="The main data image.")
507 mask: MaskSerializationModel[P] = pydantic.Field(
508 description="Bitmask that annotates the main image's pixels."
509 )
510 variance: ImageSerializationModel[P] = pydantic.Field(
511 description="Per-pixel variance estimates for the main image."
512 )
513 projection: ProjectionSerializationModel[P] | None = pydantic.Field(
514 default=None,
515 exclude_if=is_none,
516 description="Projection that maps the pixel grid to the sky.",
517 )
519 @property
520 def bbox(self) -> Box:
521 """The bounding box of the image."""
522 return self.image.bbox