Coverage for python/lsst/images/_image.py: 23%

219 statements  

« prev     ^ index     » next       coverage.py v7.14.1, created at 2026-06-10 18:00 +0000

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. 

11 

12from __future__ import annotations 

13 

14__all__ = ("Image", "ImageSerializationModel") 

15 

16from collections.abc import Callable, Sequence 

17from contextlib import ExitStack 

18from types import EllipsisType 

19from typing import TYPE_CHECKING, Any, ClassVar, final 

20 

21import astropy.io.fits 

22import astropy.units 

23import astropy.wcs 

24import numpy as np 

25import numpy.typing as npt 

26import pydantic 

27 

28from lsst.resources import ResourcePath, ResourcePathExpression 

29 

30from . import fits 

31from ._generalized_image import GeneralizedImage 

32from ._geom import YX, Box 

33from ._transforms import Frame, SkyProjection, SkyProjectionSerializationModel 

34from .serialization import ( 

35 ArchiveTree, 

36 ArrayReferenceModel, 

37 ArrayReferenceQuantityModel, 

38 InlineArrayModel, 

39 InlineArrayQuantityModel, 

40 InputArchive, 

41 InvalidParameterError, 

42 MetadataValue, 

43 OutputArchive, 

44 no_header_updates, 

45) 

46from .utils import is_none 

47 

48if TYPE_CHECKING: 

49 try: 

50 from lsst.afw.image import Image as LegacyImage 

51 except ImportError: 

52 type LegacyImage = Any # type: ignore[no-redef] 

53 

54 

55@final 

56class Image(GeneralizedImage): 

57 """A 2-d array that may be augmented with units and a nonzero origin. 

58 

59 Parameters 

60 ---------- 

61 array_or_fill 

62 Array or fill value for the image. If a fill value, ``bbox`` or 

63 ``shape`` must be provided. 

64 bbox 

65 Bounding box for the image. 

66 yx0 

67 Logical coordinates of the first pixel in the array, ordered ``y``, 

68 ``x`` (unless an `XY` instance is passed). Ignored if 

69 ``bbox`` is provided. Defaults to zeros. 

70 shape 

71 Leading dimensions of the array, ordered ``y``, ``x`` (unless an `XY` 

72 instance is passed). Only needed if ``array_or_fill`` is not an 

73 array and ``bbox`` is not provided. Like the bbox, this does not 

74 include the last dimension of the array. 

75 dtype 

76 Pixel data type override. 

77 unit 

78 Units for the image's pixel values. 

79 sky_projection 

80 Projection that maps the pixel grid to the sky. 

81 metadata 

82 Arbitrary flexible metadata to associate with the image. 

83 

84 Notes 

85 ----- 

86 Indexing the `array` attribute of an `Image` does not take into account its 

87 ``yx0`` offset, but accessing a subimage by indexing an `Image` with a 

88 `Box` does, and the `bbox` of the subimage is set to match its location 

89 within the original image. 

90 

91 Indexed assignment to a subimage requires consistency between the 

92 coordinate systems and units of both operands, but it will automatically 

93 select a subimage of the right-hand side and convert compatible units when 

94 possible. In other words:: 

95 

96 a[box] = b 

97 

98 is a shortcut for 

99 

100 a[box].quantity = b[box].quantity 

101 

102 An ellipsis (``...``) can be used instead of a `Box` to assign to the full 

103 image. 

104 """ 

105 

106 def __init__( 

107 self, 

108 array_or_fill: np.ndarray | int | float = 0, 

109 /, 

110 *, 

111 bbox: Box | None = None, 

112 yx0: Sequence[int] | None = None, 

113 shape: Sequence[int] | None = None, 

114 dtype: npt.DTypeLike | None = None, 

115 unit: astropy.units.UnitBase | None = None, 

116 sky_projection: SkyProjection[Any] | None = None, 

117 metadata: dict[str, MetadataValue] | None = None, 

118 ): 

119 super().__init__(metadata) 

120 if isinstance(array_or_fill, np.ndarray): 

121 if dtype is not None: 

122 array = np.array(array_or_fill, dtype=dtype, copy=None) 

123 else: 

124 array = array_or_fill 

125 if bbox is None: 

126 bbox = Box.from_shape(array.shape, start=yx0) 

127 elif bbox.shape != array.shape: 

128 raise ValueError( 

129 f"Explicit bbox shape {bbox.shape} does not match array with shape {array.shape}." 

130 ) 

131 if shape is not None and shape != array.shape: 

132 raise ValueError(f"Explicit shape {shape} does not match array with shape {array.shape}.") 

133 else: 

134 if bbox is None: 

135 if shape is None: 

136 raise TypeError("No bbox, shape, or array provided.") 

137 bbox = Box.from_shape(shape, start=yx0) 

138 elif shape is not None and shape != bbox.shape: 

139 raise ValueError(f"Explicit shape {shape} does not match bbox shape {bbox.shape}.") 

140 array = np.full(bbox.shape, array_or_fill, dtype=dtype) 

141 self._array: np.ndarray = array 

142 self._bbox: Box = bbox 

143 self._unit = unit 

144 self._sky_projection = sky_projection 

145 

146 @property 

147 def array(self) -> np.ndarray: 

148 """The low-level array (`numpy.ndarray`). 

149 

150 Assigning to this attribute modifies the existing array in place; the 

151 bounding box and underlying data pointer are never changed. 

152 """ 

153 return self._array 

154 

155 @array.setter 

156 def array(self, value: np.ndarray | int | float) -> None: 

157 self._array[...] = value 

158 

159 @property 

160 def quantity(self) -> astropy.units.Quantity: 

161 """The low-level array with units (`astropy.units.Quantity`). 

162 

163 Assigning to this attribute modifies the existing array in place; the 

164 bounding box and underlying data pointer are never changed. 

165 """ 

166 return astropy.units.Quantity(self._array, self._unit, copy=False) 

167 

168 @quantity.setter 

169 def quantity(self, value: astropy.units.Quantity) -> None: 

170 self.quantity[...] = value 

171 

172 @property 

173 def bbox(self) -> Box: 

174 """Bounding box for the image (`Box`).""" 

175 return self._bbox 

176 

177 @property 

178 def unit(self) -> astropy.units.UnitBase | None: 

179 """Units for the image's pixel values (`astropy.units.Unit` or 

180 `None`). 

181 """ 

182 return self._unit 

183 

184 @property 

185 def sky_projection(self) -> SkyProjection[Any] | None: 

186 """The projection that maps this image's pixel grid to the sky 

187 (`SkyProjection` | `None`). 

188 

189 Notes 

190 ----- 

191 The pixel coordinates used by this projection account for the bounding 

192 box ``start``; they are not just array indices. 

193 """ 

194 return self._sky_projection 

195 

196 def __getitem__(self, bbox: Box | EllipsisType) -> Image: 

197 if bbox is ...: 

198 return self 

199 super().__getitem__(bbox) 

200 indices = bbox.slice_within(self._bbox) 

201 return self._transfer_metadata( 

202 Image(self._array[indices], bbox=bbox, unit=self._unit, sky_projection=self._sky_projection), 

203 bbox=bbox, 

204 ) 

205 

206 def __setitem__(self, bbox: Box | EllipsisType, value: Image) -> None: 

207 self[bbox].quantity[...] = value.quantity 

208 

209 def __str__(self) -> str: 

210 return f"Image({self.bbox!s}, {self.array.dtype.type.__name__})" 

211 

212 def __repr__(self) -> str: 

213 return f"Image(..., bbox={self.bbox!r}, dtype={self.array.dtype!r})" 

214 

215 def __eq__(self, other: object) -> bool: 

216 if not isinstance(other, Image): 

217 return NotImplemented 

218 return ( 

219 self._bbox == other._bbox 

220 and self._unit == other._unit 

221 and np.array_equal(self._array, other._array, equal_nan=True) 

222 ) 

223 

224 def copy(self) -> Image: 

225 return self._transfer_metadata( 

226 Image(self._array.copy(), bbox=self._bbox, unit=self._unit, sky_projection=self._sky_projection), 

227 copy=True, 

228 ) 

229 

230 def view( 

231 self, 

232 *, 

233 unit: astropy.units.UnitBase | None | EllipsisType = ..., 

234 sky_projection: SkyProjection | None | EllipsisType = ..., 

235 yx0: Sequence[int] | EllipsisType = ..., 

236 ) -> Image: 

237 """Make a view of the image, with optional updates.""" 

238 if unit is ...: 

239 unit = self._unit 

240 if sky_projection is ...: 

241 sky_projection = self._sky_projection 

242 if yx0 is ...: 

243 yx0 = self._bbox.start 

244 return self._transfer_metadata(Image(self._array, yx0=yx0, unit=unit, sky_projection=sky_projection)) 

245 

246 def serialize[P: pydantic.BaseModel]( 

247 self, 

248 archive: OutputArchive[P], 

249 *, 

250 update_header: Callable[[astropy.io.fits.Header], None] = no_header_updates, 

251 save_projection: bool = True, 

252 add_offset_wcs: str | None = "A", 

253 tile_shape: tuple[int, ...] | None = None, 

254 options_name: str | None = None, 

255 ) -> ImageSerializationModel[P]: 

256 """Serialize the image to an output archive. 

257 

258 Parameters 

259 ---------- 

260 archive 

261 Archive to write to. 

262 update_header 

263 A callback that will be given the FITS header for the HDU 

264 containing this image in order to add keys to it. This callback 

265 may be provided but will not be called if the output format is not 

266 FITS. 

267 save_projection 

268 If `True`, save the `SkyProjection` attached to the image, if there 

269 is one. This does not affect whether a FITS WCS corresponding to 

270 the projection is written (it always is, if available, and if 

271 ``add_offset_wcs`` is not ``" "``). 

272 add_offset_wcs 

273 A FITS WCS single-character suffix to use when adding a linear 

274 WCS that maps the FITS array to the logical pixel coordinates 

275 defined by ``bbox.start``. Set to `None` to not write this WCS. 

276 If this is set to ``" "``, it will prevent the `SkyProjection` from 

277 being saved as a FITS WCS. 

278 tile_shape 

279 The recommended shape of each tile, if the archive will save 

280 the array in distinct tiles for faster subarray retrieval. 

281 This is a hint; archives are not required to use this value. 

282 options_name 

283 Use this name to look up archive options. 

284 """ 

285 

286 def _update_header(header: astropy.io.fits.Header) -> None: 

287 update_header(header) 

288 if self.unit is not None: 

289 try: 

290 header["BUNIT"] = self.unit.to_string(format="fits") 

291 except ValueError: 

292 # Units not supported by FITS; write it anyway because 

293 # the accepted units are just a recommendation in the 

294 # standard. 

295 header["BUNIT"] = self.unit.to_string() 

296 if self.sky_projection is not None and add_offset_wcs != " ": 

297 if self.fits_wcs: 

298 header.update(self.fits_wcs.to_header(relax=True)) 

299 if add_offset_wcs is not None: 

300 fits.add_offset_wcs(header, x=self.bbox.x.start, y=self.bbox.y.start, key=add_offset_wcs) 

301 

302 array_model = archive.add_array( 

303 self.array, update_header=_update_header, tile_shape=tile_shape, options_name=options_name 

304 ) 

305 serialized_projection: SkyProjectionSerializationModel[P] | None = None 

306 if save_projection and self.sky_projection is not None: 

307 serialized_projection = archive.serialize_direct("sky_projection", self.sky_projection.serialize) 

308 data = array_model if self.unit is None else array_model.with_units(self.unit) 

309 return ImageSerializationModel.model_construct( 

310 data=data, 

311 yx0=list(self.bbox.start), 

312 sky_projection=serialized_projection, 

313 metadata=self.metadata, 

314 ) 

315 

316 @staticmethod 

317 def _get_archive_tree_type[P: pydantic.BaseModel]( 

318 pointer_type: type[P], 

319 ) -> type[ImageSerializationModel[P]]: 

320 """Return the serialization model type for this object for an archive 

321 type that uses the given pointer type. 

322 """ 

323 return ImageSerializationModel[pointer_type] # type: ignore 

324 

325 _archive_default_name: ClassVar[str] = "image" 

326 """The name this object should be serialized with when written as the 

327 top-level object. 

328 """ 

329 

330 @staticmethod 

331 def from_legacy(legacy: LegacyImage, unit: astropy.units.UnitBase | None = None) -> Image: 

332 """Convert from an `lsst.afw.image.Image` instance. 

333 

334 Parameters 

335 ---------- 

336 legacy 

337 An `lsst.afw.image.Image` instance that will share pixel data with 

338 the returned object. 

339 unit 

340 Units of the image. 

341 """ 

342 return Image(legacy.array, yx0=YX(y=legacy.getY0(), x=legacy.getX0()), unit=unit) 

343 

344 def to_legacy(self, *, copy: bool | None = None) -> LegacyImage: 

345 """Convert to an `lsst.afw.image.Image` instance. 

346 

347 Parameters 

348 ---------- 

349 copy 

350 If `True`, always copy the pixel data. If `False`, return a view, 

351 and raise `TypeError` if the pixel data is read-only (this is not 

352 supported by afw). If `None`, onyl if the pixel data is 

353 read-only. 

354 """ 

355 import lsst.afw.image 

356 import lsst.geom 

357 

358 array = self._array 

359 if copy: 

360 array = array.copy() 

361 elif not self._array.flags.writeable: 

362 if copy is None: 

363 array = array.copy() 

364 else: 

365 raise TypeError("Cannot create a legacy lsst.afw.image.Image view into a read-only array.") 

366 

367 return lsst.afw.image.Image( 

368 array, 

369 deep=False, 

370 dtype=array.dtype.type, 

371 xy0=lsst.geom.Point2I(self._bbox.x.min, self._bbox.y.min), 

372 ) 

373 

374 @staticmethod 

375 def read_legacy( 

376 uri: ResourcePathExpression, 

377 *, 

378 preserve_quantization: bool = False, 

379 ext: str | int = 1, 

380 fits_wcs_frame: Frame | None = None, 

381 ) -> Image: 

382 """Read a FITS file written by `lsst.afw.image.Image.writeFits`. 

383 

384 Parameters 

385 ---------- 

386 uri 

387 URI or file name. 

388 preserve_quantization 

389 If `True`, ensure that writing the image back out again will 

390 exactly preserve quantization-compressed pixel values. This causes 

391 the arrays to be marked as read-only and stores the original binary 

392 table data for those planes in memory. If the `Image` is copied, 

393 the precompressed pixel values are not transferred to the copy. 

394 ext 

395 Name or index of the FITS HDU to read. 

396 fits_wcs_frame 

397 If not `None` and the HDU containing the image has a FITS WCS, 

398 attach a `SkyProjection` to the returned image by converting that 

399 WCS. 

400 """ 

401 opaque_metadata = fits.FitsOpaqueMetadata() 

402 with ExitStack() as exit_stack: 

403 fs, fspath = ResourcePath(uri).to_fsspec() 

404 stream = exit_stack.enter_context(fs.open(fspath)) 

405 hdu_list = exit_stack.enter_context(astropy.io.fits.open(stream)) 

406 opaque_metadata.extract_legacy_primary_header(hdu_list[0].header) 

407 bintable_hdu: astropy.io.fits.BinTableHDU | None = None 

408 if preserve_quantization: 

409 bintable_stream = exit_stack.enter_context(fs.open(fspath)) 

410 bintable_hdu_list = exit_stack.enter_context( 

411 astropy.io.fits.open(bintable_stream, disable_image_compression=True) 

412 ) 

413 bintable_hdu = bintable_hdu_list[ext] 

414 result = Image._read_legacy_hdu( 

415 hdu_list[ext], opaque_metadata, preserve_bintable=bintable_hdu, fits_wcs_frame=fits_wcs_frame 

416 ) 

417 result._opaque_metadata = opaque_metadata 

418 return result 

419 

420 @staticmethod 

421 def _read_legacy_hdu( 

422 hdu: astropy.io.fits.ImageHDU | astropy.io.fits.CompImageHDU, 

423 opaque_metadata: fits.FitsOpaqueMetadata, 

424 *, 

425 preserve_bintable: astropy.io.fits.BinTableHDU | None, 

426 fits_wcs_frame: Frame | None = None, 

427 ) -> Image: 

428 unit: astropy.units.UnitBase | None = None 

429 if (fits_unit := hdu.header.pop("BUNIT", None)) is not None: 

430 try: 

431 unit = astropy.units.Unit(fits_unit, format="fits") 

432 except ValueError: 

433 # Accept non-FITS units by assuming Astropy can still figure 

434 # them out if we don't specify the format. 

435 unit = astropy.units.Unit(fits_unit) 

436 if opaque_metadata.get_instrumental_unit() == astropy.units.electron: 

437 # Fix incorrect BUNIT='adu' in LSST preliminary_visit_image. 

438 if unit == astropy.units.adu: 

439 unit = astropy.units.electron 

440 if unit == astropy.units.adu**2: 

441 unit = astropy.units.electron**2 

442 dx: int = hdu.header.pop("LTV1") 

443 dy: int = hdu.header.pop("LTV2") 

444 yx0 = YX(y=-dy, x=-dx) 

445 read_only: bool = False 

446 if preserve_bintable is not None: 

447 opaque_metadata.precompressed[hdu.name] = fits.PrecompressedImage.from_bintable(preserve_bintable) 

448 read_only = True 

449 sky_projection: SkyProjection | None = None 

450 if fits_wcs_frame is not None: 

451 try: 

452 fits_wcs = astropy.wcs.WCS(hdu.header) 

453 except KeyError: 

454 pass 

455 else: 

456 sky_projection = SkyProjection.from_fits_wcs( 

457 fits_wcs, pixel_frame=fits_wcs_frame, x0=yx0.x, y0=yx0.y 

458 ) 

459 image = Image(hdu.data, yx0=yx0, unit=unit, sky_projection=sky_projection) 

460 if read_only: 

461 image._array.flags["WRITEABLE"] = False 

462 fits.strip_wcs_cards(hdu.header) 

463 hdu.header.strip() 

464 hdu.header.remove("EXTTYPE", ignore_missing=True) 

465 hdu.header.remove("INHERIT", ignore_missing=True) 

466 hdu.header.remove("UZSCALE", ignore_missing=True) 

467 opaque_metadata.add_header(hdu.header) 

468 return image 

469 

470 

471class ImageSerializationModel[P: pydantic.BaseModel](ArchiveTree): 

472 """Pydantic model used to represent the serialized form of an `.Image`.""" 

473 

474 SCHEMA_NAME: ClassVar[str] = "image" 

475 SCHEMA_VERSION: ClassVar[str] = "1.0.0" 

476 MIN_READ_VERSION: ClassVar[int] = 1 

477 PUBLIC_TYPE: ClassVar[type] = Image 

478 

479 data: ArrayReferenceQuantityModel | ArrayReferenceModel | InlineArrayModel | InlineArrayQuantityModel = ( 

480 pydantic.Field(description="Reference to pixel data.") 

481 ) 

482 yx0: list[int] = pydantic.Field( 

483 description="Coordinate of the first pixels in the array, ordered (y, x)." 

484 ) 

485 sky_projection: SkyProjectionSerializationModel[P] | None = pydantic.Field( 

486 default=None, 

487 exclude_if=is_none, 

488 description="Projection that maps the logical pixel grid onto the sky.", 

489 ) 

490 

491 @property 

492 def bbox(self) -> Box: 

493 """The bounding box of the image.""" 

494 match self.data: 

495 case ArrayReferenceQuantityModel() | InlineArrayQuantityModel(): 

496 shape = self.data.value.shape 

497 case ArrayReferenceModel() | InlineArrayModel(): 

498 shape = self.data.shape 

499 return Box.from_shape(shape, self.yx0) 

500 

501 def deserialize( 

502 self, 

503 archive: InputArchive[Any], 

504 *, 

505 bbox: Box | None = None, 

506 strip_header: Callable[[astropy.io.fits.Header], None] = no_header_updates, 

507 **kwargs: Any, 

508 ) -> Image: 

509 """Deserialize an image from an input archive. 

510 

511 Parameters 

512 ---------- 

513 archive 

514 Archive to read from. 

515 bbox 

516 Bounding box of a subimage to read instead. 

517 strip_header 

518 A callable that strips out any FITS header cards added by the 

519 ``update_header`` argument in the corresponding call to 

520 `Image.serialize`. 

521 **kwargs 

522 Unsupported keyword arguments are accepted only to provide better 

523 error messages (raising `serialization.InvalidParameterError`). 

524 """ 

525 if kwargs: 

526 raise InvalidParameterError(f"Unrecognized parameters for Image: {set(kwargs.keys())}.") 

527 array_model: ArrayReferenceModel | InlineArrayModel 

528 unit: astropy.units.UnitBase | None = None 

529 if isinstance(self.data, ArrayReferenceQuantityModel | InlineArrayQuantityModel): 

530 array_model = self.data.value 

531 unit = self.data.unit 

532 else: 

533 array_model = self.data 

534 

535 def _strip_header(header: astropy.io.fits.Header) -> None: 

536 if unit is not None: 

537 header.pop("BUNIT", None) 

538 fits.strip_wcs_cards(header) 

539 strip_header(header) 

540 

541 slices = bbox.slice_within(self.bbox) if bbox is not None else ... 

542 array = archive.get_array(array_model, strip_header=_strip_header, slices=slices) 

543 sky_projection = self.sky_projection.deserialize(archive) if self.sky_projection is not None else None 

544 return Image( 

545 array, 

546 yx0=self.yx0 if bbox is None else bbox.start, 

547 unit=unit, 

548 sky_projection=sky_projection, 

549 )._finish_deserialize(self) 

550 

551 def deserialize_component(self, component: str, archive: InputArchive[Any], **kwargs: Any) -> Any: 

552 if kwargs: 

553 raise InvalidParameterError(f"Unsupported parameters for Image components: {set(kwargs.keys())}.") 

554 return super().deserialize_component(component, archive)