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