Coverage for python / lsst / images / _image.py: 23%
238 statements
« prev ^ index » next coverage.py v7.13.4, created at 2026-03-01 02:11 -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__ = ("Image", "ImageSerializationModel")
16from collections.abc import Callable, Sequence
17from contextlib import ExitStack
18from functools import cached_property
19from types import EllipsisType
20from typing import Any, ClassVar, final
22import astropy.io.fits
23import astropy.units
24import astropy.wcs
25import numpy as np
26import numpy.typing as npt
27import pydantic
28from astro_metadata_translator import ObservationInfo
30from lsst.resources import ResourcePath, ResourcePathExpression
32from . import fits
33from ._geom import YX, Box
34from ._transforms import Frame, Projection, ProjectionAstropyView, ProjectionSerializationModel
35from .serialization import (
36 ArchiveTree,
37 ArrayReferenceModel,
38 ArrayReferenceQuantityModel,
39 InputArchive,
40 OpaqueArchiveMetadata,
41 OutputArchive,
42 no_header_updates,
43)
44from .utils import is_none
47@final
48class Image:
49 """A 2-d array that may be augmented with units and a nonzero origin.
51 Parameters
52 ----------
53 array_or_fill
54 Array or fill value for the image. If a fill value, ``bbox`` or
55 ``shape`` must be provided.
56 bbox
57 Bounding box for the image.
58 start
59 Logical coordinates of the first pixel in the array, ordered ``y``,
60 ``x`` (unless an `XY` instance is passed). Ignored if
61 ``bbox`` is provided. Defaults to zeros.
62 shape
63 Leading dimensions of the array, ordered ``y``, ``x`` (unless an `XY`
64 instance is passed). Only needed if ``array_or_fill`` is not an
65 array and ``bbox`` is not provided. Like the bbox, this does not
66 include the last dimension of the array.
67 dtype
68 Pixel data type override.
69 unit
70 Units for the image's pixel values.
71 projection
72 Projection that maps the pixel grid to the sky.
73 obs_info
74 General information about the associated observation in standardized
75 form.
77 Notes
78 -----
79 Indexing the `array` attribute of an `Image` does not take into account its
80 ``start`` offset, but accessing a subimage by indexing an `Image` with a
81 `Box` does, and the `bbox` of the subimage is set to match its location
82 within the original image.
84 Indexed assignment to a subimage requires consistency between the
85 coordinate systems and units of both operands, but it will automatically
86 select a subimage of the right-hand side and convert compatible units when
87 possible. In other words::
89 a[box] = b
91 is a shortcut for
93 a[box].quantity = b[box].quantity
95 An ellipsis (``...``) can be used instead of a `Box` to assign to the full
96 image.
97 """
99 def __init__(
100 self,
101 array_or_fill: np.ndarray | int | float = 0,
102 /,
103 *,
104 bbox: Box | None = None,
105 start: Sequence[int] | None = None,
106 shape: Sequence[int] | None = None,
107 dtype: npt.DTypeLike | None = None,
108 unit: astropy.units.UnitBase | None = None,
109 projection: Projection[Any] | None = None,
110 obs_info: ObservationInfo | None = None,
111 ):
112 if isinstance(array_or_fill, np.ndarray):
113 if dtype is not None:
114 array = np.array(array_or_fill, dtype=dtype)
115 else:
116 array = array_or_fill
117 if bbox is None:
118 bbox = Box.from_shape(array.shape, start=start)
119 elif bbox.shape != array.shape:
120 raise ValueError(
121 f"Explicit bbox shape {bbox.shape} does not match array with shape {array.shape}."
122 )
123 if shape is not None and shape != array.shape:
124 raise ValueError(f"Explicit shape {shape} does not match array with shape {array.shape}.")
125 else:
126 if bbox is None:
127 if shape is None:
128 raise TypeError("No bbox, shape, or array provided.")
129 bbox = Box.from_shape(shape, start=start)
130 elif shape is not None and shape != bbox.shape:
131 raise ValueError(f"Explicit shape {shape} does not match bbox shape {bbox.shape}.")
132 array = np.full(bbox.shape, array_or_fill, dtype=dtype)
133 self._array: np.ndarray = array
134 self._bbox: Box = bbox
135 self._unit = unit
136 self._projection = projection
137 self._opaque_metadata: OpaqueArchiveMetadata | None = None
138 self._obs_info = obs_info
140 @property
141 def array(self) -> np.ndarray:
142 """The low-level array (`numpy.ndarray`).
144 Assigning to this attribute modifies the existing array in place; the
145 bounding box and underlying data pointer are never changed.
146 """
147 return self._array
149 @array.setter
150 def array(self, value: np.ndarray | int | float) -> None:
151 self._array[...] = value
153 @property
154 def quantity(self) -> astropy.units.Quantity:
155 """The low-level array with units (`astropy.units.Quantity`).
157 Assigning to this attribute modifies the existing array in place; the
158 bounding box and underlying data pointer are never changed.
159 """
160 return astropy.units.Quantity(self._array, self._unit, copy=False)
162 @quantity.setter
163 def quantity(self, value: astropy.units.Quantity) -> None:
164 self.quantity[...] = value
166 @property
167 def bbox(self) -> Box:
168 """Bounding box for the image (`Box`)."""
169 return self._bbox
171 @property
172 def unit(self) -> astropy.units.UnitBase | None:
173 """Units for the image's pixel values (`astropy.units.Unit` or
174 `None`).
175 """
176 return self._unit
178 @property
179 def projection(self) -> Projection[Any] | None:
180 """The projection that maps this image's pixel grid to the sky
181 (`Projection` | `None`).
183 Notes
184 -----
185 The pixel coordinates used by this projection account for the bounding
186 box ``start``; they are not just array indices.
187 """
188 return self._projection
190 @property
191 def obs_info(self) -> ObservationInfo | None:
192 """General information about the associated observation in standard
193 form. (`~astro_metadata_translator.ObservationInfo` | `None`).
194 """
195 return self._obs_info
197 @property
198 def astropy_wcs(self) -> ProjectionAstropyView | None:
199 """An Astropy WCS for this image's pixel array.
201 Notes
202 -----
203 As expected for Astropy WCS objects, this defines pixel coordinates
204 such that the first row and column in `array` are ``(0, 0)``, not
205 ``bbox.start``, as is the case for `projection`.
207 This object satisfies the `astropy.wcs.wcsapi.BaseHighLevelWCS` and
208 `astropy.wcs.wcsapi.BaseLowLevelWCS` interfaces, but it is not an
209 `astropy.wcs.WCS` (use `fits_wcs` for that).
210 """
211 return self._projection.as_astropy(self.bbox) if self._projection is not None else None
213 @cached_property
214 def fits_wcs(self) -> astropy.wcs.WCS | None:
215 """An Astropy FITS WCS for this image's pixel array.
217 Notes
218 -----
219 As expected for Astropy WCS objects, this defines pixel coordinates
220 such that the first row and column in `array` are ``(0, 0)``, not
221 ``bbox.start``, as is the case for `projection`.
223 This may be an approximation or absent if `projection` is not
224 naturally representable as a FITS WCS.
225 """
226 return (
227 self._projection.as_fits_wcs(self.bbox, allow_approximation=True)
228 if self._projection is not None
229 else None
230 )
232 def __getitem__(self, bbox: Box | EllipsisType | tuple[slice, slice]) -> Image:
233 indices: EllipsisType | tuple[slice, ...]
234 if bbox is ...:
235 indices = ...
236 bbox = self._bbox
237 elif isinstance(bbox, tuple):
238 bbox = self._bbox[bbox]
239 indices = bbox.slice_within(self._bbox)
240 else:
241 indices = bbox.slice_within(self._bbox)
242 result = Image(
243 self._array[indices],
244 bbox=bbox,
245 unit=self._unit,
246 projection=self._projection,
247 obs_info=self._obs_info,
248 )
249 if self._opaque_metadata is not None:
250 result._opaque_metadata = self._opaque_metadata.subset(bbox)
251 return result
253 def __setitem__(self, bbox: Box | EllipsisType, value: Image) -> None:
254 if bbox is ...:
255 bbox = self._bbox
256 if value._bbox != bbox:
257 value = value[bbox]
258 # Use the quantity property to handle unit conversions and checks.
259 self.quantity[bbox.slice_within(self._bbox)] = value.quantity
261 def __str__(self) -> str:
262 return f"Image({self.bbox!s}, {self.array.dtype.type.__name__})"
264 def __repr__(self) -> str:
265 return f"Image(..., bbox={self.bbox!r}, dtype={self.array.dtype!r})"
267 def __eq__(self, other: object) -> bool:
268 if not isinstance(other, Image):
269 return NotImplemented
270 return (
271 self._bbox == other._bbox
272 and self._unit == other._unit
273 and np.array_equal(self._array, other._array, equal_nan=True)
274 )
276 def copy(self) -> Image:
277 """Deep-copy the image, with optional updates."""
278 result = Image(
279 self._array.copy(),
280 bbox=self._bbox,
281 unit=self._unit,
282 projection=self._projection,
283 obs_info=self._obs_info,
284 )
285 if self._opaque_metadata is not None:
286 result._opaque_metadata = self._opaque_metadata.copy()
287 return result
289 def view(
290 self,
291 *,
292 unit: astropy.units.UnitBase | None | EllipsisType = ...,
293 projection: Projection | None | EllipsisType = ...,
294 start: Sequence[int] | EllipsisType = ...,
295 obs_info: ObservationInfo | None | EllipsisType = ...,
296 ) -> Image:
297 """Make a view of the image, with optional updates."""
298 if unit is ...:
299 unit = self._unit
300 if projection is ...:
301 projection = self._projection
302 if start is ...:
303 start = self._bbox.start
304 if obs_info is ...:
305 obs_info = self._obs_info
306 result = Image(self._array, start=start, unit=unit, projection=projection, obs_info=obs_info)
307 if self._opaque_metadata is not None:
308 result._opaque_metadata = self._opaque_metadata.copy()
309 return result
311 def serialize[P: pydantic.BaseModel](
312 self,
313 archive: OutputArchive[P],
314 *,
315 update_header: Callable[[astropy.io.fits.Header], None] = no_header_updates,
316 save_projection: bool = True,
317 add_offset_wcs: str | None = "A",
318 ) -> ImageSerializationModel[P]:
319 """Serialize the image to an output archive.
321 Parameters
322 ----------
323 archive
324 Archive to write to.
325 update_header
326 A callback that will be given the FITS header for the HDU
327 containing this image in order to add keys to it. This callback
328 may be provided but will not be called if the output format is not
329 FITS.
330 save_projection
331 If `True`, save the `Projection` attached to the image, if there
332 is one.
333 add_offset_wcs
334 A FITS WCS single-character suffix to use when adding a linear
335 WCS that maps the FITS array to the logical pixel coordinates
336 defined by ``bbox.start``. Set to `None` to not write this WCS.
337 """
338 if save_projection and add_offset_wcs == "":
339 raise TypeError("save_projection=True is not compatible with add_offset_wcs=''.")
341 def _update_header(header: astropy.io.fits.Header) -> None:
342 update_header(header)
343 if self.unit is not None:
344 header["BUNIT"] = self.unit.to_string(format="fits")
345 if self.projection is not None:
346 if self.fits_wcs:
347 header.update(self.fits_wcs.to_header(relax=True))
348 if add_offset_wcs is not None:
349 fits.add_offset_wcs(header, x=self.bbox.x.start, y=self.bbox.y.start, key=add_offset_wcs)
351 ref = archive.add_array(self.array, update_header=_update_header)
352 serialized_projection: ProjectionSerializationModel[P] | None = None
353 if save_projection and self.projection is not None:
354 serialized_projection = archive.serialize_direct("projection", self.projection.serialize)
355 if self.unit is None:
356 return ImageSerializationModel.model_construct(
357 data=ref,
358 start=list(self.bbox.start),
359 projection=serialized_projection,
360 obs_info=self._obs_info,
361 )
362 else:
363 return ImageSerializationModel.model_construct(
364 data=ArrayReferenceQuantityModel.model_construct(value=ref, unit=self.unit),
365 start=list(self.bbox.start),
366 projection=serialized_projection,
367 obs_info=self._obs_info,
368 )
370 @staticmethod
371 def deserialize(
372 model: ImageSerializationModel[Any],
373 archive: InputArchive[Any],
374 *,
375 bbox: Box | None = None,
376 strip_header: Callable[[astropy.io.fits.Header], None] = no_header_updates,
377 ) -> Image:
378 """Deserialize an image from an input archive.
380 Parameters
381 ----------
382 model
383 A Pydantic model representation of the image, holding references
384 to data stored in the archive.
385 archive
386 Archive to read from.
387 bbox
388 Bounding box of a subimage to read instead.
389 strip_header
390 A callable that strips out any FITS header cards added by the
391 ``update_header`` argument in the corresponding call to
392 `serialize`.
393 """
394 ref: ArrayReferenceModel
395 unit: astropy.units.UnitBase | None = None
396 if isinstance(model.data, ArrayReferenceQuantityModel):
397 ref = model.data.value
398 unit = model.data.unit
399 else:
400 ref = model.data
402 def _strip_header(header: astropy.io.fits.Header) -> None:
403 if unit is not None:
404 header.pop("BUNIT", None)
405 fits.strip_wcs_cards(header)
406 strip_header(header)
408 slices = bbox.slice_within(model.bbox) if bbox is not None else ...
409 array = archive.get_array(ref, strip_header=_strip_header, slices=slices)
410 projection = (
411 Projection.deserialize(model.projection, archive) if model.projection is not None else None
412 )
413 return Image(
414 array,
415 start=model.start if bbox is None else bbox.start,
416 unit=unit,
417 projection=projection,
418 obs_info=model.obs_info,
419 )
421 @staticmethod
422 def _get_archive_tree_type[P: pydantic.BaseModel](
423 pointer_type: type[P],
424 ) -> type[ImageSerializationModel[P]]:
425 """Return the serialization model type for this object for an archive
426 type that uses the given pointer type.
427 """
428 return ImageSerializationModel[pointer_type] # type: ignore
430 _archive_default_name: ClassVar[str] = "image"
431 """The name this object should be serialized with when written as the
432 top-level object.
433 """
435 def write_fits(
436 self,
437 filename: str,
438 *,
439 compression: fits.FitsCompressionOptions | None = fits.FitsCompressionOptions.DEFAULT,
440 compression_seed: int | None = None,
441 ) -> None:
442 """Write the image to a FITS file.
444 Parameters
445 ----------
446 filename
447 Name of the file to write to. Must be a local file.
448 compression
449 Compression options.
450 compression_seed
451 A FITS tile compression seed to use whenever the configured
452 compression seed is `None` or (for backwards compatibility) ``0``.
453 """
454 compression_options = {}
455 if compression is not fits.FitsCompressionOptions.DEFAULT:
456 compression_options[self._archive_default_name] = compression
457 fits.write(self, filename, compression_options, compression_seed=compression_seed)
459 @staticmethod
460 def read_fits(url: ResourcePathExpression, *, bbox: Box | None = None) -> Image:
461 """Read an image from a FITS file.
463 Parameters
464 ----------
465 url
466 URL of the file to read; may be any type supported by
467 `lsst.resources.ResourcePath`.
468 bbox
469 Bounding box of a subimage to read instead.
470 """
471 return fits.read(Image, url, bbox=bbox)
473 @staticmethod
474 def from_legacy(legacy: Any, unit: astropy.units.Unit | None = None) -> Image:
475 """Convert from an `lsst.afw.image.Image` instance.
477 Parameters
478 ----------
479 legacy
480 An `lsst.afw.image.Image` instance that will share pixel data with
481 the returned object.
482 unit
483 Units of the image.
484 """
485 return Image(legacy.array, start=(legacy.getY0(), legacy.getX0()), unit=unit)
487 def to_legacy(self, *, copy: bool | None = None) -> Any:
488 """Convert to an `lsst.afw.image.Image` instance.
490 Parameters
491 ----------
492 copy
493 If `True`, always copy the pixel data. If `False`, return a view,
494 and raise `TypeError` if the pixel data is read-only (this is not
495 supported by afw). If `None`, onyl if the pixel data is
496 read-only.
497 """
498 import lsst.afw.image
499 import lsst.geom
501 array = self._array
502 if copy:
503 array = array.copy()
504 elif not self._array.flags.writeable:
505 if copy is None:
506 array = array.copy()
507 else:
508 raise TypeError("Cannot create a legacy lsst.afw.image.Image view into a read-only array.")
510 return lsst.afw.image.Image(
511 array,
512 deep=False,
513 dtype=array.dtype.type,
514 xy0=lsst.geom.Point2I(self._bbox.x.min, self._bbox.y.min),
515 )
517 @staticmethod
518 def read_legacy(
519 uri: ResourcePathExpression,
520 *,
521 preserve_quantization: bool = False,
522 ext: str | int = 1,
523 fits_wcs_frame: Frame | None = None,
524 ) -> Image:
525 """Read a FITS file written by `lsst.afw.image.Image.writeFits`.
527 Parameters
528 ----------
529 uri
530 URI or file name.
531 preserve_quantization
532 If `True`, ensure that writing the image back out again will
533 exactly preserve quantization-compressed pixel values. This causes
534 the arrays to be marked as read-only and stores the original binary
535 table data for those planes in memory. If the `Image` is copied,
536 the precompressed pixel values are not transferred to the copy.
537 ext
538 Name or index of the FITS HDU to read.
539 fits_wcs_frame
540 If not `None` and the HDU containing the image has a FITS WCS,
541 attach a `Projection` to the returned image by converting that
542 WCS.
543 """
544 opaque_metadata = fits.FitsOpaqueMetadata()
545 with ExitStack() as exit_stack:
546 fs, fspath = ResourcePath(uri).to_fsspec()
547 stream = exit_stack.enter_context(fs.open(fspath))
548 hdu_list = exit_stack.enter_context(astropy.io.fits.open(stream))
549 opaque_metadata.extract_legacy_primary_header(hdu_list[0].header)
550 bintable_hdu: astropy.io.fits.BinTableHDU | None = None
551 if preserve_quantization:
552 bintable_stream = exit_stack.enter_context(fs.open(fspath))
553 bintable_hdu_list = exit_stack.enter_context(
554 astropy.io.fits.open(bintable_stream, disable_image_compression=True)
555 )
556 bintable_hdu = bintable_hdu_list[ext]
557 result = Image._read_legacy_hdu(
558 hdu_list[ext], opaque_metadata, preserve_bintable=bintable_hdu, fits_wcs_frame=fits_wcs_frame
559 )
560 result._opaque_metadata = opaque_metadata
561 return result
563 @staticmethod
564 def _read_legacy_hdu(
565 hdu: astropy.io.fits.ImageHDU | astropy.io.fits.CompImageHDU,
566 opaque_metadata: fits.FitsOpaqueMetadata,
567 *,
568 preserve_bintable: astropy.io.fits.BinTableHDU | None,
569 fits_wcs_frame: Frame | None = None,
570 ) -> Image:
571 unit: astropy.units.UnitBase | None = None
572 if (fits_unit := hdu.header.pop("BUNIT", None)) is not None:
573 unit = astropy.units.Unit(fits_unit, format="fits")
574 dx: int = hdu.header.pop("LTV1")
575 dy: int = hdu.header.pop("LTV2")
576 start = YX(y=-dy, x=-dx)
577 read_only: bool = False
578 if preserve_bintable is not None:
579 opaque_metadata.precompressed[hdu.name] = fits.PrecompressedImage.from_bintable(preserve_bintable)
580 read_only = True
581 projection: Projection | None = None
582 if fits_wcs_frame is not None:
583 try:
584 fits_wcs = astropy.wcs.WCS(hdu.header)
585 except KeyError:
586 pass
587 else:
588 projection = Projection.from_fits_wcs(
589 fits_wcs, pixel_frame=fits_wcs_frame, x0=start.x, y0=start.y
590 )
591 image = Image(hdu.data, start=start, unit=unit, projection=projection)
592 if read_only:
593 image._array.flags["WRITEABLE"] = False
594 fits.strip_wcs_cards(hdu.header)
595 hdu.header.strip()
596 hdu.header.remove("EXTTYPE", ignore_missing=True)
597 hdu.header.remove("INHERIT", ignore_missing=True)
598 hdu.header.remove("UZSCALE", ignore_missing=True)
599 opaque_metadata.add_header(hdu.header)
600 return image
603class ImageSerializationModel[P: pydantic.BaseModel](ArchiveTree):
604 """Pydantic model used to represent the serialized form of an `.Image`."""
606 data: ArrayReferenceQuantityModel | ArrayReferenceModel = pydantic.Field(
607 description="Reference to pixel data."
608 )
609 start: list[int] = pydantic.Field(
610 description="Coordinate of the first pixels in the array, ordered (y, x)."
611 )
612 projection: ProjectionSerializationModel[P] | None = pydantic.Field(
613 default=None,
614 exclude_if=is_none,
615 description="Projection that maps the logical pixel grid onto the sky.",
616 )
617 obs_info: ObservationInfo | None = pydantic.Field(
618 default=None,
619 exclude_if=is_none,
620 description="Standardized description of image metadata",
621 )
623 @property
624 def bbox(self) -> Box:
625 """The bounding box of the image."""
626 if isinstance(self.data, ArrayReferenceQuantityModel):
627 shape = self.data.value.shape
628 else:
629 shape = self.data.shape
630 return Box.from_shape(shape, self.start)