Coverage for python/lsst/summit/extras/animation.py: 12%
223 statements
« prev ^ index » next coverage.py v6.5.0, created at 2023-01-06 03:38 -0800
« prev ^ index » next coverage.py v6.5.0, created at 2023-01-06 03:38 -0800
1# This file is part of summit_extras.
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/>.
22import os
23import subprocess
24import shutil
25import uuid
26import math
28import matplotlib.pyplot as plt
30from lsst.pipe.tasks.quickFrameMeasurement import QuickFrameMeasurementTask, QuickFrameMeasurementTaskConfig
31from lsst.atmospec.utils import getTargetCentroidFromWcs
32import lsst.afw.display as afwDisplay
33import lsst.afw.math as afwMath
34import logging
35import lsst.meas.algorithms as measAlg
36from lsst.summit.utils.utils import dayObsIntToString, setupLogging
37from lsst.summit.utils.butlerUtils import (datasetExists, getExpRecordFromDataId, makeDefaultLatissButler,
38 getDayObs, getSeqNum, updateDataIdOrDataCord,
39 getLatissOnSkyDataIds)
41from lsst.atmospec.utils import airMassFromRawMetadata
42logger = logging.getLogger("lsst.summit.extras.animation")
43setupLogging()
46class Animator():
47 """Animate the list of dataIds in the order in which they are specified
48 for the data product specified."""
50 def __init__(self, butler, dataIdList, outputPath, outputFilename, *,
51 remakePngs=False,
52 clobberVideoAndGif=False,
53 keepIntermediateGif=False,
54 smoothImages=True,
55 plotObjectCentroids=True,
56 useQfmForCentroids=False,
57 dataProductToPlot='calexp',
58 debug=False):
60 self.butler = butler
61 self.dataIdList = dataIdList
62 self.outputPath = outputPath
63 self.outputFilename = os.path.join(outputPath, outputFilename)
64 if not self.outputFilename.endswith(".mp4"):
65 self.outputFilename += ".mp4"
66 self.pngPath = os.path.join(outputPath, "pngs/")
68 self.remakePngs = remakePngs
69 self.clobberVideoAndGif = clobberVideoAndGif
70 self.keepIntermediateGif = keepIntermediateGif
71 self.smoothImages = smoothImages
72 self.plotObjectCentroids = plotObjectCentroids
73 self.useQfmForCentroids = useQfmForCentroids
74 self.dataProductToPlot = dataProductToPlot
75 self.debug = debug
77 # zfilled at the start as animation is alphabetical
78 # if you're doing more than 1e6 files you've got bigger problems
79 self.toAnimateTemplate = "%06d-%s-%s.png"
80 self.basicTemplate = "%s-%s.png"
82 qfmTaskConfig = QuickFrameMeasurementTaskConfig()
83 self.qfmTask = QuickFrameMeasurementTask(config=qfmTaskConfig)
85 afwDisplay.setDefaultBackend("matplotlib")
86 self.fig = plt.figure(figsize=(15, 15))
87 self.disp = afwDisplay.Display(self.fig)
88 self.disp.setImageColormap('gray')
89 self.disp.scale('asinh', 'zscale')
91 self.pngsToMakeDataIds = []
92 self.preRun() # sets the above list
94 @staticmethod
95 def _strDataId(dataId):
96 """Make a dataId into a string suitable for use as a filename.
98 Parameters
99 ----------
100 dataId : `dict`
101 The data id.
103 Returns
104 -------
105 strId : `str`
106 The data id as a string.
107 """
108 if (dayObs := getDayObs(dataId)) and (seqNum := getSeqNum(dataId)): # nicely ordered if easy
109 return f"{dayObsIntToString(dayObs)}-{seqNum:05d}"
111 # General case (and yeah, I should probably learn regex someday)
112 dIdStr = str(dataId)
113 dIdStr = dIdStr.replace(' ', "")
114 dIdStr = dIdStr.replace('{', "")
115 dIdStr = dIdStr.replace('}', "")
116 dIdStr = dIdStr.replace('\'', "")
117 dIdStr = dIdStr.replace(':', "-")
118 dIdStr = dIdStr.replace(',', "-")
119 return dIdStr
121 def dataIdToFilename(self, dataId, includeNumber=False, imNum=None):
122 """Convert dataId to filename.
124 Returns a full path+filename by default. if includeNumber then
125 returns just the filename for use in temporary dir for animation."""
126 if includeNumber:
127 assert imNum is not None
129 dIdStr = self._strDataId(dataId)
131 if includeNumber: # for use in temp dir, so not full path
132 filename = self.toAnimateTemplate%(imNum, dIdStr, self.dataProductToPlot)
133 return os.path.join(filename)
134 else:
135 filename = self.basicTemplate%(dIdStr, self.dataProductToPlot)
136 return os.path.join(self.pngPath, filename)
138 def exists(self, obj):
139 if type(obj) == str:
140 return os.path.exists(obj)
141 raise RuntimeError("Other type checks not yet implemented")
143 def preRun(self):
144 # check the paths work
145 if not os.path.exists(self.pngPath):
146 os.makedirs(self.pngPath)
147 assert os.path.exists(self.pngPath), f"Failed to create output dir: {self.pngsPath}"
149 if self.exists(self.outputFilename):
150 if self.clobberVideoAndGif:
151 os.remove(self.outputFilename)
152 else:
153 raise RuntimeError(f"Output file {self.outputFilename} exists and clobber==False")
155 # make list of found & missing files
156 dIdsWithPngs = [d for d in self.dataIdList if self.exists(self.dataIdToFilename(d))]
157 dIdsWithoutPngs = [d for d in self.dataIdList if d not in dIdsWithPngs]
158 if self.debug:
159 logger.info(f"dIdsWithPngs = {dIdsWithPngs}")
160 logger.info(f"dIdsWithoutPngs = {dIdsWithoutPngs}")
162 # check the datasets exist for the pngs which need remaking
163 missingData = [d for d in dIdsWithoutPngs if not datasetExists(self.butler, self.dataProductToPlot, d,
164 detector=0)]
166 logger.info(f"Of the provided {len(self.dataIdList)} dataIds:")
167 logger.info(f"{len(dIdsWithPngs)} existing pngs were found")
168 logger.info(f"{len(dIdsWithoutPngs)} do not yet exist")
170 if missingData:
171 for dId in missingData:
172 msg = f"Failed to find {self.dataProductToPlot} for {dId}"
173 logger.warning(msg)
174 self.dataIdList.remove(dId)
175 logger.info(f"Of the {len(dIdsWithoutPngs)} dataIds without pngs, {len(missingData)}"
176 " did not have the corresponding dataset existing")
178 if self.remakePngs:
179 self.pngsToMakeDataIds = [d for d in self.dataIdList if d not in missingData]
180 else:
181 self.pngsToMakeDataIds = [d for d in dIdsWithoutPngs if d not in missingData]
183 msg = f"So {len(self.pngsToMakeDataIds)} will be made"
184 if self.remakePngs and len(dIdsWithPngs) > 0:
185 msg += " because remakePngs=True"
186 logger.info(msg)
188 def run(self):
189 # make the missing pngs
190 if self.pngsToMakeDataIds:
191 logger.info('Creating necessary pngs...')
192 for i, dataId in enumerate(self.pngsToMakeDataIds):
193 logger.info(f'Making png for file {i+1} of {len(self.pngsToMakeDataIds)}')
194 self.makePng(dataId, self.dataIdToFilename(dataId))
196 # stage files in temp dir with numbers prepended to filenames
197 if not self.dataIdList:
198 logger.warning('No files to animate - nothing to do')
199 return
201 logger.info('Copying files to ordered temp dir...')
202 pngFilesOriginal = [self.dataIdToFilename(d) for d in self.dataIdList]
203 for filename in pngFilesOriginal: # these must all now exist, but let's assert just in case
204 assert self.exists(filename)
205 tempDir = os.path.join(self.pngPath, f"{uuid.uuid1()}/"[0:8])
206 os.makedirs(tempDir)
207 pngFileList = [] # list of number-prepended files in the temp dir
208 for i, dId in enumerate(self.dataIdList):
209 srcFile = self.dataIdToFilename(dId)
210 destFile = os.path.join(tempDir, self.dataIdToFilename(dId, includeNumber=True, imNum=i))
211 shutil.copy(srcFile, destFile)
212 pngFileList.append(destFile)
214 # # create gif in temp dir
215 # outputGifFilename = os.path.join(tempDir, 'animation.gif')
216 # self.pngsToGif(pngFileList, outputGifFilename)
218 # # gif turn into mp4, optionally keep gif by moving up to output dir
219 # logger.info('Turning gif into mp4...')
220 # outputMp4Filename = self.outputFilename
221 # self.gifToMp4(outputGifFilename, outputMp4Filename)
223 # self.tidyUp(tempDir)
224 # logger.info('Finished!')
226 # create gif in temp dir
228 logger.info('Making mp4 of pngs...')
229 self.pngsToMp4(tempDir, self.outputFilename, 10, verbose=False)
230 self.tidyUp(tempDir)
231 logger.info(f'Finished! Output at {self.outputFilename}')
232 return self.outputFilename
234 def _titleFromExp(self, exp, dataId):
235 expRecord = getExpRecordFromDataId(self.butler, dataId)
236 obj = expRecord.target_name
237 expTime = expRecord.exposure_time
238 filterCompound = expRecord.physical_filter
239 filt, grating = filterCompound.split('~')
240 rawMd = self.butler.get('raw.metadata', dataId)
241 airmass = airMassFromRawMetadata(rawMd) # XXX this could be improved a lot
242 if not airmass:
243 airmass = -1
244 dayObs = dayObsIntToString(getDayObs(dataId))
245 timestamp = expRecord.timespan.begin.to_datetime().strftime("%H:%M:%S") # no microseconds
246 ms = expRecord.timespan.begin.to_datetime().strftime("%f") # always 6 chars long, 000000 if zero
247 timestamp += f".{ms[0:2]}"
248 title = f"seqNum {getSeqNum(dataId)} - {dayObs} {timestamp}TAI - "
249 title += f"Object: {obj} expTime: {expTime}s Filter: {filt} Grating: {grating} Airmass: {airmass:.3f}"
250 return title
252 def getStarPixCoord(self, exp, doMotionCorrection=True, useQfm=False):
253 target = exp.visitInfo.object
255 if self.useQfmForCentroids:
256 try:
257 result = self.qfmTask.run(exp)
258 pixCoord = result.brightestObjCentroid
259 expId = exp.info.id
260 logger.info(f'expId {expId} has centroid {pixCoord}')
261 except Exception:
262 return None
263 else:
264 pixCoord = getTargetCentroidFromWcs(exp, target, doMotionCorrection=doMotionCorrection)
265 return pixCoord
267 def makePng(self, dataId, saveFilename):
268 if self.exists(saveFilename) and not self.remakePngs: # should not be possible due to prerun
269 assert False, f"Almost overwrote {saveFilename} - how is this possible?"
271 if self.debug:
272 logger.info(f"Creating {saveFilename}")
274 self.fig.clear()
276 # must always keep exp unsmoothed for the centroiding via qfm
277 try:
278 exp = self.butler.get(self.dataProductToPlot, dataId)
279 except Exception:
280 # oh no, that should never happen, but it does! Let's just skip
281 logger.warning(f'Skipped {dataId}, because {self.dataProductToPlot} retrieval failed!')
282 return
283 toDisplay = exp
284 if self.smoothImages:
285 toDisplay = exp.clone()
286 toDisplay = self._smoothExp(toDisplay, 2)
288 try:
289 self.disp.mtv(toDisplay.image, title=self._titleFromExp(exp, dataId))
290 self.disp.scale('asinh', 'zscale')
291 except RuntimeError: # all-nan images slip through and don't display
292 self.disp.scale('linear', 0, 1)
293 self.disp.mtv(toDisplay.image, title=self._titleFromExp(exp, dataId))
294 self.disp.scale('asinh', 'zscale') # set back for next image
295 pass
297 if self.plotObjectCentroids:
298 try:
299 pixCoord = self.getStarPixCoord(exp)
300 if pixCoord:
301 self.disp.dot('x', *pixCoord, ctype='C1', size=50)
302 self.disp.dot('o', *pixCoord, ctype='C1', size=50)
303 else:
304 self.disp.dot('x', 2000, 2000, ctype='red', size=2000)
305 except Exception:
306 logger.warning(f"Failed to find OBJECT location for {dataId}")
308 deltaH = -0.05
309 deltaV = -0.05
310 plt.subplots_adjust(right=1+deltaH, left=0-deltaH, top=1+deltaV, bottom=0-deltaV)
311 self.fig.savefig(saveFilename)
312 logger.info(f'Saved png for {dataId} to {saveFilename}')
314 def pngsToMp4(self, indir, outfile, framerate, verbose=False):
315 """Create the movie with ffmpeg, from files."""
316 # NOTE: the order of ffmpeg arguments *REALLY MATTERS*.
317 # Reorder them at your own peril!
318 pathPattern = f'\"{os.path.join(indir, "*.png")}\"'
319 if verbose:
320 ffmpeg_verbose = 'info'
321 else:
322 ffmpeg_verbose = 'error'
323 cmd = ['ffmpeg',
324 '-v', ffmpeg_verbose,
325 '-f', 'image2',
326 '-y',
327 '-pattern_type glob',
328 '-framerate', f'{framerate}',
329 '-i', pathPattern,
330 '-vcodec', 'libx264',
331 '-b:v', '20000k',
332 '-profile:v', 'main',
333 '-pix_fmt', 'yuv420p',
334 '-threads', '10',
335 '-r', f'{framerate}',
336 os.path.join(outfile)]
338 subprocess.check_call(r' '.join(cmd), shell=True)
340 def tidyUp(self, tempDir):
341 shutil.rmtree(tempDir)
342 return
344 def _smoothExp(self, exp, smoothing, kernelSize=7):
345 """Use for DISPLAY ONLY!
347 Return a smoothed copy of the exposure
348 with the original mask plane in place."""
349 psf = measAlg.DoubleGaussianPsf(kernelSize, kernelSize, smoothing/(2*math.sqrt(2*math.log(2))))
350 newExp = exp.clone()
351 originalMask = exp.mask
353 kernel = psf.getKernel()
354 afwMath.convolve(newExp.maskedImage, newExp.maskedImage, kernel, afwMath.ConvolutionControl())
355 newExp.mask = originalMask
356 return newExp
359def animateDay(butler, dayObs, outputPath, dataProductToPlot='quickLookExp'):
360 outputFilename = f'{dayObs}.mp4'
362 onSkyIds = getLatissOnSkyDataIds(butler, startDate=dayObs, endDate=dayObs)
363 logger.info(f"Found {len(onSkyIds)} on sky ids for {dayObs}")
365 onSkyIds = [updateDataIdOrDataCord(dataId, detector=0) for dataId in onSkyIds]
367 animator = Animator(butler, onSkyIds, outputPath, outputFilename,
368 dataProductToPlot=dataProductToPlot,
369 remakePngs=False,
370 debug=False,
371 clobberVideoAndGif=True,
372 plotObjectCentroids=True,
373 useQfmForCentroids=True)
374 filename = animator.run()
375 return filename
378if __name__ == '__main__': 378 ↛ 379line 378 didn't jump to line 379, because the condition on line 378 was never true
379 outputPath = '/home/mfl/animatorOutput/main/'
380 butler = makeDefaultLatissButler()
382 day = 20211104
383 animateDay(butler, day, outputPath)