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

225 statements  

« prev     ^ index     » next       coverage.py v7.13.4, created at 2026-03-03 00:49 -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. 

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 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 

27from astro_metadata_translator import ObservationInfo 

28 

29from lsst.resources import ResourcePath, ResourcePathExpression 

30 

31from . import fits 

32from ._generalized_image import GeneralizedImage 

33from ._geom import YX, Box 

34from ._transforms import Frame, Projection, 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 

45 

46 

47@final 

48class Image(GeneralizedImage): 

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

50 

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. 

76 

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. 

83 

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:: 

88 

89 a[box] = b 

90 

91 is a shortcut for 

92 

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

94 

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

96 image. 

97 """ 

98 

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 

139 

140 @property 

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

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

143 

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 

148 

149 @array.setter 

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

151 self._array[...] = value 

152 

153 @property 

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

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

156 

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) 

161 

162 @quantity.setter 

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

164 self.quantity[...] = value 

165 

166 @property 

167 def bbox(self) -> Box: 

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

169 return self._bbox 

170 

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 

177 

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`). 

182 

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 

189 

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 

196 

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

198 super().__getitem__(bbox) 

199 if bbox is ...: 

200 return self 

201 indices = bbox.slice_within(self._bbox) 

202 result = Image( 

203 self._array[indices], 

204 bbox=bbox, 

205 unit=self._unit, 

206 projection=self._projection, 

207 obs_info=self._obs_info, 

208 ) 

209 if self._opaque_metadata is not None: 

210 result._opaque_metadata = self._opaque_metadata.subset(bbox) 

211 return result 

212 

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

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

215 

216 def __str__(self) -> str: 

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

218 

219 def __repr__(self) -> str: 

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

221 

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

223 if not isinstance(other, Image): 

224 return NotImplemented 

225 return ( 

226 self._bbox == other._bbox 

227 and self._unit == other._unit 

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

229 ) 

230 

231 def copy(self) -> Image: 

232 """Deep-copy the image, with optional updates.""" 

233 result = Image( 

234 self._array.copy(), 

235 bbox=self._bbox, 

236 unit=self._unit, 

237 projection=self._projection, 

238 obs_info=self._obs_info, 

239 ) 

240 if self._opaque_metadata is not None: 

241 result._opaque_metadata = self._opaque_metadata.copy() 

242 return result 

243 

244 def view( 

245 self, 

246 *, 

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

248 projection: Projection | None | EllipsisType = ..., 

249 start: Sequence[int] | EllipsisType = ..., 

250 obs_info: ObservationInfo | None | EllipsisType = ..., 

251 ) -> Image: 

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

253 if unit is ...: 

254 unit = self._unit 

255 if projection is ...: 

256 projection = self._projection 

257 if start is ...: 

258 start = self._bbox.start 

259 if obs_info is ...: 

260 obs_info = self._obs_info 

261 result = Image(self._array, start=start, unit=unit, projection=projection, obs_info=obs_info) 

262 if self._opaque_metadata is not None: 

263 result._opaque_metadata = self._opaque_metadata.copy() 

264 return result 

265 

266 def serialize[P: pydantic.BaseModel]( 

267 self, 

268 archive: OutputArchive[P], 

269 *, 

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

271 save_projection: bool = True, 

272 add_offset_wcs: str | None = "A", 

273 ) -> ImageSerializationModel[P]: 

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

275 

276 Parameters 

277 ---------- 

278 archive 

279 Archive to write to. 

280 update_header 

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

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

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

284 FITS. 

285 save_projection 

286 If `True`, save the `Projection` attached to the image, if there 

287 is one. 

288 add_offset_wcs 

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

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

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

292 """ 

293 if save_projection and add_offset_wcs == "": 

294 raise TypeError("save_projection=True is not compatible with add_offset_wcs=''.") 

295 

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

297 update_header(header) 

298 if self.unit is not None: 

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

300 if self.projection is not None: 

301 if self.fits_wcs: 

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

303 if add_offset_wcs is not None: 

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

305 

306 ref = archive.add_array(self.array, update_header=_update_header) 

307 serialized_projection: ProjectionSerializationModel[P] | None = None 

308 if save_projection and self.projection is not None: 

309 serialized_projection = archive.serialize_direct("projection", self.projection.serialize) 

310 if self.unit is None: 

311 return ImageSerializationModel.model_construct( 

312 data=ref, 

313 start=list(self.bbox.start), 

314 projection=serialized_projection, 

315 obs_info=self._obs_info, 

316 ) 

317 else: 

318 return ImageSerializationModel.model_construct( 

319 data=ArrayReferenceQuantityModel.model_construct(value=ref, unit=self.unit), 

320 start=list(self.bbox.start), 

321 projection=serialized_projection, 

322 obs_info=self._obs_info, 

323 ) 

324 

325 @staticmethod 

326 def deserialize( 

327 model: ImageSerializationModel[Any], 

328 archive: InputArchive[Any], 

329 *, 

330 bbox: Box | None = None, 

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

332 ) -> Image: 

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

334 

335 Parameters 

336 ---------- 

337 model 

338 A Pydantic model representation of the image, holding references 

339 to data stored in the archive. 

340 archive 

341 Archive to read from. 

342 bbox 

343 Bounding box of a subimage to read instead. 

344 strip_header 

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

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

347 `serialize`. 

348 """ 

349 ref: ArrayReferenceModel 

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

351 if isinstance(model.data, ArrayReferenceQuantityModel): 

352 ref = model.data.value 

353 unit = model.data.unit 

354 else: 

355 ref = model.data 

356 

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

358 if unit is not None: 

359 header.pop("BUNIT", None) 

360 fits.strip_wcs_cards(header) 

361 strip_header(header) 

362 

363 slices = bbox.slice_within(model.bbox) if bbox is not None else ... 

364 array = archive.get_array(ref, strip_header=_strip_header, slices=slices) 

365 projection = ( 

366 Projection.deserialize(model.projection, archive) if model.projection is not None else None 

367 ) 

368 return Image( 

369 array, 

370 start=model.start if bbox is None else bbox.start, 

371 unit=unit, 

372 projection=projection, 

373 obs_info=model.obs_info, 

374 ) 

375 

376 @staticmethod 

377 def _get_archive_tree_type[P: pydantic.BaseModel]( 

378 pointer_type: type[P], 

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

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

381 type that uses the given pointer type. 

382 """ 

383 return ImageSerializationModel[pointer_type] # type: ignore 

384 

385 _archive_default_name: ClassVar[str] = "image" 

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

387 top-level object. 

388 """ 

389 

390 def write_fits( 

391 self, 

392 filename: str, 

393 *, 

394 compression: fits.FitsCompressionOptions | None = fits.FitsCompressionOptions.DEFAULT, 

395 compression_seed: int | None = None, 

396 ) -> None: 

397 """Write the image to a FITS file. 

398 

399 Parameters 

400 ---------- 

401 filename 

402 Name of the file to write to. Must be a local file. 

403 compression 

404 Compression options. 

405 compression_seed 

406 A FITS tile compression seed to use whenever the configured 

407 compression seed is `None` or (for backwards compatibility) ``0``. 

408 """ 

409 compression_options = {} 

410 if compression is not fits.FitsCompressionOptions.DEFAULT: 

411 compression_options[self._archive_default_name] = compression 

412 fits.write(self, filename, compression_options, compression_seed=compression_seed) 

413 

414 @staticmethod 

415 def read_fits(url: ResourcePathExpression, *, bbox: Box | None = None) -> Image: 

416 """Read an image from a FITS file. 

417 

418 Parameters 

419 ---------- 

420 url 

421 URL of the file to read; may be any type supported by 

422 `lsst.resources.ResourcePath`. 

423 bbox 

424 Bounding box of a subimage to read instead. 

425 """ 

426 return fits.read(Image, url, bbox=bbox) 

427 

428 @staticmethod 

429 def from_legacy(legacy: Any, unit: astropy.units.Unit | None = None) -> Image: 

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

431 

432 Parameters 

433 ---------- 

434 legacy 

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

436 the returned object. 

437 unit 

438 Units of the image. 

439 """ 

440 return Image(legacy.array, start=(legacy.getY0(), legacy.getX0()), unit=unit) 

441 

442 def to_legacy(self, *, copy: bool | None = None) -> Any: 

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

444 

445 Parameters 

446 ---------- 

447 copy 

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

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

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

451 read-only. 

452 """ 

453 import lsst.afw.image 

454 import lsst.geom 

455 

456 array = self._array 

457 if copy: 

458 array = array.copy() 

459 elif not self._array.flags.writeable: 

460 if copy is None: 

461 array = array.copy() 

462 else: 

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

464 

465 return lsst.afw.image.Image( 

466 array, 

467 deep=False, 

468 dtype=array.dtype.type, 

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

470 ) 

471 

472 @staticmethod 

473 def read_legacy( 

474 uri: ResourcePathExpression, 

475 *, 

476 preserve_quantization: bool = False, 

477 ext: str | int = 1, 

478 fits_wcs_frame: Frame | None = None, 

479 ) -> Image: 

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

481 

482 Parameters 

483 ---------- 

484 uri 

485 URI or file name. 

486 preserve_quantization 

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

488 exactly preserve quantization-compressed pixel values. This causes 

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

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

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

492 ext 

493 Name or index of the FITS HDU to read. 

494 fits_wcs_frame 

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

496 attach a `Projection` to the returned image by converting that 

497 WCS. 

498 """ 

499 opaque_metadata = fits.FitsOpaqueMetadata() 

500 with ExitStack() as exit_stack: 

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

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

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

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

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

506 if preserve_quantization: 

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

508 bintable_hdu_list = exit_stack.enter_context( 

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

510 ) 

511 bintable_hdu = bintable_hdu_list[ext] 

512 result = Image._read_legacy_hdu( 

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

514 ) 

515 result._opaque_metadata = opaque_metadata 

516 return result 

517 

518 @staticmethod 

519 def _read_legacy_hdu( 

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

521 opaque_metadata: fits.FitsOpaqueMetadata, 

522 *, 

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

524 fits_wcs_frame: Frame | None = None, 

525 ) -> Image: 

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

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

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

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

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

531 start = YX(y=-dy, x=-dx) 

532 read_only: bool = False 

533 if preserve_bintable is not None: 

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

535 read_only = True 

536 projection: Projection | None = None 

537 if fits_wcs_frame is not None: 

538 try: 

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

540 except KeyError: 

541 pass 

542 else: 

543 projection = Projection.from_fits_wcs( 

544 fits_wcs, pixel_frame=fits_wcs_frame, x0=start.x, y0=start.y 

545 ) 

546 image = Image(hdu.data, start=start, unit=unit, projection=projection) 

547 if read_only: 

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

549 fits.strip_wcs_cards(hdu.header) 

550 hdu.header.strip() 

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

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

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

554 opaque_metadata.add_header(hdu.header) 

555 return image 

556 

557 

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

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

560 

561 data: ArrayReferenceQuantityModel | ArrayReferenceModel = pydantic.Field( 

562 description="Reference to pixel data." 

563 ) 

564 start: list[int] = pydantic.Field( 

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

566 ) 

567 projection: ProjectionSerializationModel[P] | None = pydantic.Field( 

568 default=None, 

569 exclude_if=is_none, 

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

571 ) 

572 obs_info: ObservationInfo | None = pydantic.Field( 

573 default=None, 

574 exclude_if=is_none, 

575 description="Standardized description of image metadata", 

576 ) 

577 

578 @property 

579 def bbox(self) -> Box: 

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

581 if isinstance(self.data, ArrayReferenceQuantityModel): 

582 shape = self.data.value.shape 

583 else: 

584 shape = self.data.shape 

585 return Box.from_shape(shape, self.start)