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