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

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
38def writeFits(filename, stamp_ims, metadata, write_mask, write_variance):
39 """Write a single FITS file containing all stamps.
41 Parameters
42 ----------
43 filename : `str`
44 A string indicating the output filename
45 stamps_ims : iterable of `lsst.afw.image.MaskedImageF`
46 An iterable of masked images
47 metadata : `PropertyList`
48 A collection of key, value metadata pairs to be
49 written to the primary header
50 write_mask : `bool`
51 Write the mask data to the output file?
52 write_variance : `bool`
53 Write the variance data to the output file?
54 """
55 metadata['HAS_MASK'] = write_mask
56 metadata['HAS_VARIANCE'] = write_variance
57 metadata['N_STAMPS'] = len(stamp_ims)
58 # create primary HDU with global metadata
59 fitsPrimary = afwFits.Fits(filename, "w")
60 fitsPrimary.createEmpty()
61 fitsPrimary.writeMetadata(metadata)
62 fitsPrimary.closeFile()
64 # add all pixel data optionally writing mask and variance information
65 for i, stamp in enumerate(stamp_ims):
66 metadata = PropertyList()
67 # EXTVER should be 1-based, the index from enumerate is 0-based
68 metadata.update({'EXTVER': i+1, 'EXTNAME': 'IMAGE'})
69 stamp.getImage().writeFits(filename, metadata=metadata, mode='a')
70 if write_mask:
71 metadata = PropertyList()
72 metadata.update({'EXTVER': i+1, 'EXTNAME': 'MASK'})
73 stamp.getMask().writeFits(filename, metadata=metadata, mode='a')
74 if write_variance:
75 metadata = PropertyList()
76 metadata.update({'EXTVER': i+1, 'EXTNAME': 'VARIANCE'})
77 stamp.getVariance().writeFits(filename, metadata=metadata, mode='a')
78 return None
81def readFitsWithOptions(filename, stamp_factory, options):
82 """Read stamps from FITS file, allowing for only a
83 subregion of the stamps to be read.
85 Parameters
86 ----------
87 filename : `str`
88 A string indicating the file to read
89 stamp_factory : classmethod
90 A factory function defined on a dataclass for constructing
91 stamp objects a la `lsst.meas.alrogithm.Stamp`
92 options : `PropertyList` or `dict`
93 A collection of parameters. If it contains a bounding box
94 (``bbox`` key), or if certain other keys (``llcX``, ``llcY``,
95 ``width``, ``height``) are available for one to be constructed,
96 the bounding box is passed to the ``FitsReader`` in order to
97 return a sub-image.
99 Returns
100 -------
101 stamps : `list` of dataclass objects like `Stamp`, PropertyList
102 A tuple of a list of `Stamp`-like objects
103 metadata : `PropertyList`
104 The metadata
105 """
106 # extract necessary info from metadata
107 metadata = afwFits.readMetadata(filename, hdu=0)
108 f = afwFits.Fits(filename, 'r')
109 nExtensions = f.countHdus()
110 nStamps = metadata["N_STAMPS"]
111 # check if a bbox was provided
112 kwargs = {}
113 if options:
114 # gen3 API
115 if "bbox" in options.keys():
116 kwargs["bbox"] = options["bbox"]
117 # gen2 API
118 elif "llcX" in options.keys():
119 llcX = options["llcX"]
120 llcY = options["llcY"]
121 width = options["width"]
122 height = options["height"]
123 bbox = Box2I(Point2I(llcX, llcY), Extent2I(width, height))
124 kwargs["bbox"] = bbox
125 stamp_parts = {}
126 # We need to be careful because nExtensions includes the primary
127 # header data unit
128 for idx in range(nExtensions-1):
129 md = afwFits.readMetadata(filename, hdu=idx+1)
130 if md['EXTNAME'] in ('IMAGE', 'VARIANCE'):
131 reader = afwImage.ImageFitsReader(filename, hdu=idx+1)
132 elif md['EXTNAME'] == 'MASK':
133 reader = afwImage.MaskFitsReader(filename, hdu=idx+1)
134 else:
135 raise ValueError(f"Unknown extension type: {md['EXTNAME']}")
136 stamp_parts.setdefault(md['EXTVER'], {})[md['EXTNAME'].lower()] = reader.read(**kwargs)
137 if len(stamp_parts) != nStamps:
138 raise ValueError(f'Number of stamps read ({len(stamp_parts)}) does not agree with the '
139 f'number of stamps recorded in the metadata ({nStamps}).')
140 # construct stamps themselves
141 stamps = []
142 for k in range(nStamps):
143 # Need to increment by one since EXTVER starts at 1
144 maskedImage = afwImage.MaskedImageF(**stamp_parts[k+1])
145 stamps.append(stamp_factory(maskedImage, metadata, k))
147 return stamps, metadata
150@dataclass
151class AbstractStamp(abc.ABC):
152 """Single abstract stamp
154 Parameters
155 ----------
156 Inherit from this class to add metadata to the stamp
157 """
159 @classmethod
160 @abc.abstractmethod
161 def factory(cls, stamp_im, metadata, index):
162 """This method is needed to service the FITS reader.
163 We need a standard interface to construct objects like this.
164 Parameters needed to construct this object are passed in via
165 a metadata dictionary and then passed to the constructor of
166 this class.
168 Parameters
169 ----------
170 stamp : `lsst.afw.image.MaskedImage`
171 Pixel data to pass to the constructor
172 metadata : `dict`
173 Dictionary containing the information
174 needed by the constructor.
175 idx : `int`
176 Index into the lists in ``metadata``
178 Returns
179 -------
180 stamp : `AbstractStamp`
181 An instance of this class
182 """
183 raise NotImplementedError
186@dataclass
187class Stamp(AbstractStamp):
188 """Single stamp
190 Parameters
191 ----------
192 stamp_im : `lsst.afw.image.MaskedImageF`
193 The actual pixel values for the postage stamp
194 position : `lsst.geom.SpherePoint`
195 Position of the center of the stamp. Note the user
196 must keep track of the coordinate system
197 """
198 stamp_im: afwImage.maskedImage.MaskedImageF
199 position: SpherePoint
201 @classmethod
202 def factory(cls, stamp_im, metadata, index):
203 """This method is needed to service the FITS reader.
204 We need a standard interface to construct objects like this.
205 Parameters needed to construct this object are passed in via
206 a metadata dictionary and then passed to the constructor of
207 this class. If lists of values are passed with the following
208 keys, they will be passed to the constructor, otherwise dummy
209 values will be passed: RA_DEG, DEC_DEG. They should
210 each point to lists of values.
212 Parameters
213 ----------
214 stamp : `lsst.afw.image.MaskedImage`
215 Pixel data to pass to the constructor
216 metadata : `dict`
217 Dictionary containing the information
218 needed by the constructor.
219 idx : `int`
220 Index into the lists in ``metadata``
222 Returns
223 -------
224 stamp : `Stamp`
225 An instance of this class
226 """
227 if 'RA_DEG' in metadata and 'DEC_DEG' in metadata:
228 return cls(stamp_im=stamp_im,
229 position=SpherePoint(Angle(metadata.getArray('RA_DEG')[index], degrees),
230 Angle(metadata.getArray('DEC_DEG')[index], degrees)))
231 else:
232 return cls(stamp_im=stamp_im, position=SpherePoint(Angle(numpy.nan), Angle(numpy.nan)))
235class StampsBase(abc.ABC, Sequence):
236 """Collection of stamps and associated metadata.
238 Parameters
239 ----------
240 stamps : iterable
241 This should be an iterable of dataclass objects
242 a la ``lsst.meas.algorithms.Stamp``.
243 metadata : `lsst.daf.base.PropertyList`, optional
244 Metadata associated with the bright stars.
245 use_mask : `bool`, optional
246 If ``True`` read and write the mask data. Default ``True``.
247 use_variance : `bool`, optional
248 If ``True`` read and write the variance data. Default ``True``.
250 Notes
251 -----
252 A butler can be used to read only a part of the stamps,
253 specified by a bbox:
255 >>> starSubregions = butler.get("brightStarStamps_sub", dataId, bbox=bbox)
256 """
258 def __init__(self, stamps, metadata=None, use_mask=True, use_variance=True):
259 if not hasattr(stamps, '__iter__'):
260 raise ValueError('The stamps parameter must be iterable.')
261 for stamp in stamps:
262 if not isinstance(stamp, AbstractStamp):
263 raise ValueError('The entries in stamps must inherit from AbstractStamp. '
264 f'Got {type(stamp)}.')
265 self._stamps = stamps
266 self._metadata = PropertyList() if metadata is None else metadata.deepCopy()
267 self.use_mask = use_mask
268 self.use_variance = use_variance
270 @classmethod
271 @abc.abstractmethod
272 def readFits(cls, filename):
273 """Build an instance of this class from a file.
275 Parameters
276 ----------
277 filename : `str`
278 Name of the file to read
279 """
280 raise NotImplementedError
282 @classmethod
283 @abc.abstractmethod
284 def readFitsWithOptions(cls, filename, options):
285 """Build an instance of this class with options.
287 Parameters
288 ----------
289 filename : `str`
290 Name of the file to read
291 options : `PropertyList`
292 Collection of metadata parameters
293 """
294 raise NotImplementedError
296 @abc.abstractmethod
297 def _refresh_metadata(self):
298 """Make sure metadata is up to date since this object
299 can be extende
300 """
301 raise NotImplementedError
303 def writeFits(self, filename):
304 """Write this object to a file.
306 Parameters
307 ----------
308 filename : `str`
309 Name of file to write
310 """
311 self._refresh_metadata()
312 stamps_ims = self.getMaskedImages()
313 writeFits(filename, stamps_ims, self._metadata, self.use_mask, self.use_variance)
315 def __len__(self):
316 return len(self._stamps)
318 def __getitem__(self, index):
319 return self._stamps[index]
321 def __iter__(self):
322 return iter(self._stamps)
324 def getMaskedImages(self):
325 """Retrieve star images.
327 Returns
328 -------
329 maskedImages :
330 `list` [`lsst.afw.image.maskedImage.maskedImage.MaskedImageF`]
331 """
332 return [stamp.stamp_im for stamp in self._stamps]
334 @property
335 def metadata(self):
336 return self._metadata
339class Stamps(StampsBase):
340 def _refresh_metadata(self):
341 positions = self.getPositions()
342 self._metadata['RA_DEG'] = [p.getRa().asDegrees() for p in positions]
343 self._metadata['DEC_DEG'] = [p.getDec().asDegrees() for p in positions]
345 def getPositions(self):
346 return [s.position for s in self._stamps]
348 def append(self, item):
349 """Add an additional stamp.
351 Parameters
352 ----------
353 item : `Stamp`
354 Stamp object to append.
355 """
356 if not isinstance(item, Stamp):
357 raise ValueError("Ojbects added must be a Stamp object.")
358 self._stamps.append(item)
359 return None
361 def extend(self, stamp_list):
362 """Extend Stamps instance by appending elements from another instance.
364 Parameters
365 ----------
366 stamps_list : `list` [`Stamp`]
367 List of Stamp object to append.
368 """
369 for s in stamp_list:
370 if not isinstance(s, Stamp):
371 raise ValueError('Can only extend with Stamp objects')
372 self._stamps += stamp_list
374 @classmethod
375 def readFits(cls, filename):
376 """Build an instance of this class from a file.
378 Parameters
379 ----------
380 filename : `str`
381 Name of the file to read
383 Returns
384 -------
385 object : `Stamps`
386 An instance of this class
387 """
388 return cls.readFitsWithOptions(filename, None)
390 @classmethod
391 def readFitsWithOptions(cls, filename, options):
392 """Build an instance of this class with options.
394 Parameters
395 ----------
396 filename : `str`
397 Name of the file to read
398 options : `PropertyList` or `dict`
399 Collection of metadata parameters
401 Returns
402 -------
403 object : `Stamps`
404 An instance of this class
405 """
406 stamps, metadata = readFitsWithOptions(filename, Stamp.factory, options)
407 return cls(stamps, metadata=metadata, use_mask=metadata['HAS_MASK'],
408 use_variance=metadata['HAS_VARIANCE'])