Coverage for python/lsst/meas/algorithms/stamps.py : 26%

Hot-keys on this page
r m x p toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1# This file is part of meas_algorithms.
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# This program is free software: you can redistribute it and/or modify
10# it under the terms of the GNU General Public License as published by
11# the Free Software Foundation, either version 3 of the License, or
12# (at your option) any later version.
13#
14# This program is distributed in the hope that it will be useful,
15# but WITHOUT ANY WARRANTY; without even the implied warranty of
16# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
17# GNU General Public License for more details.
18#
19# You should have received a copy of the GNU General Public License
20# along with this program. If not, see <https://www.gnu.org/licenses/>.
21#
22"""Collection of small images (stamps).
23"""
25__all__ = ["Stamp", "Stamps", "StampsBase", "writeFits", "readFitsWithOptions"]
27from collections.abc import Sequence
28import abc
29from dataclasses import dataclass
30import numpy
32import lsst.afw.image as afwImage
33import lsst.afw.fits as afwFits
34from lsst.geom import Box2I, Point2I, Extent2I, Angle, degrees, SpherePoint
35from lsst.daf.base import PropertyList
36from lsst.daf.butler.core.utils import getFullTypeName
37from lsst.utils import doImport
40def writeFits(filename, stamp_ims, metadata, type_name, write_mask, write_variance):
41 """Write a single FITS file containing all stamps.
43 Parameters
44 ----------
45 filename : `str`
46 A string indicating the output filename
47 stamps_ims : iterable of `lsst.afw.image.MaskedImageF`
48 An iterable of masked images
49 metadata : `PropertyList`
50 A collection of key, value metadata pairs to be
51 written to the primary header
52 type_name : `str`
53 Python type name of the StampsBase subclass to use
54 write_mask : `bool`
55 Write the mask data to the output file?
56 write_variance : `bool`
57 Write the variance data to the output file?
58 """
59 metadata['HAS_MASK'] = write_mask
60 metadata['HAS_VARIANCE'] = write_variance
61 metadata['N_STAMPS'] = len(stamp_ims)
62 metadata['STAMPCLS'] = type_name
63 # Record version number in case of future code changes
64 metadata['VERSION'] = 1
65 # create primary HDU with global metadata
66 fitsPrimary = afwFits.Fits(filename, "w")
67 fitsPrimary.createEmpty()
68 fitsPrimary.writeMetadata(metadata)
69 fitsPrimary.closeFile()
71 # add all pixel data optionally writing mask and variance information
72 for i, stamp in enumerate(stamp_ims):
73 metadata = PropertyList()
74 # EXTVER should be 1-based, the index from enumerate is 0-based
75 metadata.update({'EXTVER': i+1, 'EXTNAME': 'IMAGE'})
76 stamp.getImage().writeFits(filename, metadata=metadata, mode='a')
77 if write_mask:
78 metadata = PropertyList()
79 metadata.update({'EXTVER': i+1, 'EXTNAME': 'MASK'})
80 stamp.getMask().writeFits(filename, metadata=metadata, mode='a')
81 if write_variance:
82 metadata = PropertyList()
83 metadata.update({'EXTVER': i+1, 'EXTNAME': 'VARIANCE'})
84 stamp.getVariance().writeFits(filename, metadata=metadata, mode='a')
85 return None
88def readFitsWithOptions(filename, stamp_factory, options):
89 """Read stamps from FITS file, allowing for only a
90 subregion of the stamps to be read.
92 Parameters
93 ----------
94 filename : `str`
95 A string indicating the file to read
96 stamp_factory : classmethod
97 A factory function defined on a dataclass for constructing
98 stamp objects a la `lsst.meas.alrogithm.Stamp`
99 options : `PropertyList` or `dict`
100 A collection of parameters. If it contains a bounding box
101 (``bbox`` key), or if certain other keys (``llcX``, ``llcY``,
102 ``width``, ``height``) are available for one to be constructed,
103 the bounding box is passed to the ``FitsReader`` in order to
104 return a sub-image.
106 Returns
107 -------
108 stamps : `list` of dataclass objects like `Stamp`, PropertyList
109 A tuple of a list of `Stamp`-like objects
110 metadata : `PropertyList`
111 The metadata
112 """
113 # extract necessary info from metadata
114 metadata = afwFits.readMetadata(filename, hdu=0)
115 f = afwFits.Fits(filename, 'r')
116 nExtensions = f.countHdus()
117 nStamps = metadata["N_STAMPS"]
118 # check if a bbox was provided
119 kwargs = {}
120 if options:
121 # gen3 API
122 if "bbox" in options.keys():
123 kwargs["bbox"] = options["bbox"]
124 # gen2 API
125 elif "llcX" in options.keys():
126 llcX = options["llcX"]
127 llcY = options["llcY"]
128 width = options["width"]
129 height = options["height"]
130 bbox = Box2I(Point2I(llcX, llcY), Extent2I(width, height))
131 kwargs["bbox"] = bbox
132 stamp_parts = {}
133 # We need to be careful because nExtensions includes the primary
134 # header data unit
135 for idx in range(nExtensions-1):
136 md = afwFits.readMetadata(filename, hdu=idx+1)
137 if md['EXTNAME'] in ('IMAGE', 'VARIANCE'):
138 reader = afwImage.ImageFitsReader(filename, hdu=idx+1)
139 elif md['EXTNAME'] == 'MASK':
140 reader = afwImage.MaskFitsReader(filename, hdu=idx+1)
141 else:
142 raise ValueError(f"Unknown extension type: {md['EXTNAME']}")
143 stamp_parts.setdefault(md['EXTVER'], {})[md['EXTNAME'].lower()] = reader.read(**kwargs)
144 if len(stamp_parts) != nStamps:
145 raise ValueError(f'Number of stamps read ({len(stamp_parts)}) does not agree with the '
146 f'number of stamps recorded in the metadata ({nStamps}).')
147 # construct stamps themselves
148 stamps = []
149 for k in range(nStamps):
150 # Need to increment by one since EXTVER starts at 1
151 maskedImage = afwImage.MaskedImageF(**stamp_parts[k+1])
152 stamps.append(stamp_factory(maskedImage, metadata, k))
154 return stamps, metadata
157@dataclass
158class AbstractStamp(abc.ABC):
159 """Single abstract stamp
161 Parameters
162 ----------
163 Inherit from this class to add metadata to the stamp
164 """
166 @classmethod
167 @abc.abstractmethod
168 def factory(cls, stamp_im, metadata, index):
169 """This method is needed to service the FITS reader.
170 We need a standard interface to construct objects like this.
171 Parameters needed to construct this object are passed in via
172 a metadata dictionary and then passed to the constructor of
173 this class.
175 Parameters
176 ----------
177 stamp : `lsst.afw.image.MaskedImage`
178 Pixel data to pass to the constructor
179 metadata : `dict`
180 Dictionary containing the information
181 needed by the constructor.
182 idx : `int`
183 Index into the lists in ``metadata``
185 Returns
186 -------
187 stamp : `AbstractStamp`
188 An instance of this class
189 """
190 raise NotImplementedError
193@dataclass
194class Stamp(AbstractStamp):
195 """Single stamp
197 Parameters
198 ----------
199 stamp_im : `lsst.afw.image.MaskedImageF`
200 The actual pixel values for the postage stamp
201 position : `lsst.geom.SpherePoint`
202 Position of the center of the stamp. Note the user
203 must keep track of the coordinate system
204 """
205 stamp_im: afwImage.maskedImage.MaskedImageF
206 position: SpherePoint
208 @classmethod
209 def factory(cls, stamp_im, metadata, index):
210 """This method is needed to service the FITS reader.
211 We need a standard interface to construct objects like this.
212 Parameters needed to construct this object are passed in via
213 a metadata dictionary and then passed to the constructor of
214 this class. If lists of values are passed with the following
215 keys, they will be passed to the constructor, otherwise dummy
216 values will be passed: RA_DEG, DEC_DEG. They should
217 each point to lists of values.
219 Parameters
220 ----------
221 stamp : `lsst.afw.image.MaskedImage`
222 Pixel data to pass to the constructor
223 metadata : `dict`
224 Dictionary containing the information
225 needed by the constructor.
226 idx : `int`
227 Index into the lists in ``metadata``
229 Returns
230 -------
231 stamp : `Stamp`
232 An instance of this class
233 """
234 if 'RA_DEG' in metadata and 'DEC_DEG' in metadata:
235 return cls(stamp_im=stamp_im,
236 position=SpherePoint(Angle(metadata.getArray('RA_DEG')[index], degrees),
237 Angle(metadata.getArray('DEC_DEG')[index], degrees)))
238 else:
239 return cls(stamp_im=stamp_im, position=SpherePoint(Angle(numpy.nan), Angle(numpy.nan)))
242class StampsBase(abc.ABC, Sequence):
243 """Collection of stamps and associated metadata.
245 Parameters
246 ----------
247 stamps : iterable
248 This should be an iterable of dataclass objects
249 a la ``lsst.meas.algorithms.Stamp``.
250 metadata : `lsst.daf.base.PropertyList`, optional
251 Metadata associated with the bright stars.
252 use_mask : `bool`, optional
253 If ``True`` read and write the mask data. Default ``True``.
254 use_variance : `bool`, optional
255 If ``True`` read and write the variance data. Default ``True``.
257 Notes
258 -----
259 A butler can be used to read only a part of the stamps,
260 specified by a bbox:
262 >>> starSubregions = butler.get("brightStarStamps_sub", dataId, bbox=bbox)
263 """
265 def __init__(self, stamps, metadata=None, use_mask=True, use_variance=True):
266 if not hasattr(stamps, '__iter__'):
267 raise ValueError('The stamps parameter must be iterable.')
268 for stamp in stamps:
269 if not isinstance(stamp, AbstractStamp):
270 raise ValueError('The entries in stamps must inherit from AbstractStamp. '
271 f'Got {type(stamp)}.')
272 self._stamps = stamps
273 self._metadata = PropertyList() if metadata is None else metadata.deepCopy()
274 self.use_mask = use_mask
275 self.use_variance = use_variance
277 @classmethod
278 def readFits(cls, filename):
279 """Build an instance of this class from a file.
281 Parameters
282 ----------
283 filename : `str`
284 Name of the file to read
285 """
287 return cls.readFitsWithOptions(filename, None)
289 @classmethod
290 def readFitsWithOptions(cls, filename, options):
291 """Build an instance of this class with options.
293 Parameters
294 ----------
295 filename : `str`
296 Name of the file to read
297 options : `PropertyList`
298 Collection of metadata parameters
299 """
300 # To avoid problems since this is no longer an abstract method
301 if cls is not StampsBase:
302 raise NotImplementedError(
303 f"Please implement specific FITS reader for class {cls}"
304 )
306 # Load metadata to get class
307 metadata = afwFits.readMetadata(filename, hdu=0)
308 type_name = metadata.get("STAMPCLS")
309 if type_name is None:
310 raise RuntimeError(
311 f"No class name in file {filename}. Unable to instantiate correct"
312 " stamps subclass. Is this an old version format Stamps file?"
313 )
315 # Import class and override `cls`
316 stamp_type = doImport(type_name)
317 cls = stamp_type
319 return cls.readFitsWithOptions(filename, options)
321 @abc.abstractmethod
322 def _refresh_metadata(self):
323 """Make sure metadata is up to date since this object
324 can be extende
325 """
326 raise NotImplementedError
328 def writeFits(self, filename):
329 """Write this object to a file.
331 Parameters
332 ----------
333 filename : `str`
334 Name of file to write
335 """
336 self._refresh_metadata()
337 stamps_ims = self.getMaskedImages()
338 type_name = getFullTypeName(self)
339 writeFits(filename, stamps_ims, self._metadata, type_name, self.use_mask, self.use_variance)
341 def __len__(self):
342 return len(self._stamps)
344 def __getitem__(self, index):
345 return self._stamps[index]
347 def __iter__(self):
348 return iter(self._stamps)
350 def getMaskedImages(self):
351 """Retrieve star images.
353 Returns
354 -------
355 maskedImages :
356 `list` [`lsst.afw.image.maskedImage.maskedImage.MaskedImageF`]
357 """
358 return [stamp.stamp_im for stamp in self._stamps]
360 @property
361 def metadata(self):
362 return self._metadata
365class Stamps(StampsBase):
366 def _refresh_metadata(self):
367 positions = self.getPositions()
368 self._metadata['RA_DEG'] = [p.getRa().asDegrees() for p in positions]
369 self._metadata['DEC_DEG'] = [p.getDec().asDegrees() for p in positions]
371 def getPositions(self):
372 return [s.position for s in self._stamps]
374 def append(self, item):
375 """Add an additional stamp.
377 Parameters
378 ----------
379 item : `Stamp`
380 Stamp object to append.
381 """
382 if not isinstance(item, Stamp):
383 raise ValueError("Objects added must be a Stamp object.")
384 self._stamps.append(item)
385 return None
387 def extend(self, stamp_list):
388 """Extend Stamps instance by appending elements from another instance.
390 Parameters
391 ----------
392 stamps_list : `list` [`Stamp`]
393 List of Stamp object to append.
394 """
395 for s in stamp_list:
396 if not isinstance(s, Stamp):
397 raise ValueError('Can only extend with Stamp objects')
398 self._stamps += stamp_list
400 @classmethod
401 def readFits(cls, filename):
402 """Build an instance of this class from a file.
404 Parameters
405 ----------
406 filename : `str`
407 Name of the file to read
409 Returns
410 -------
411 object : `Stamps`
412 An instance of this class
413 """
414 return cls.readFitsWithOptions(filename, None)
416 @classmethod
417 def readFitsWithOptions(cls, filename, options):
418 """Build an instance of this class with options.
420 Parameters
421 ----------
422 filename : `str`
423 Name of the file to read
424 options : `PropertyList` or `dict`
425 Collection of metadata parameters
427 Returns
428 -------
429 object : `Stamps`
430 An instance of this class
431 """
432 stamps, metadata = readFitsWithOptions(filename, Stamp.factory, options)
433 return cls(stamps, metadata=metadata, use_mask=metadata['HAS_MASK'],
434 use_variance=metadata['HAS_VARIANCE'])