Coverage for python/lsst/summit/extras/animation.py: 12%
221 statements
« prev ^ index » next coverage.py v6.5.0, created at 2022-11-10 03:42 -0800
« prev ^ index » next coverage.py v6.5.0, created at 2022-11-10 03:42 -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 dayObs = dayObsIntToString(getDayObs(dataId))
243 timestamp = expRecord.timespan.begin.to_datetime().strftime("%H:%M:%S") # no microseconds
244 ms = expRecord.timespan.begin.to_datetime().strftime("%f") # always 6 chars long, 000000 if zero
245 timestamp += f".{ms[0:2]}"
246 title = f"seqNum {getSeqNum(dataId)} - {dayObs} {timestamp}TAI - "
247 title += f"Object: {obj} expTime: {expTime}s Filter: {filt} Grating: {grating} Airmass: {airmass:.3f}"
248 return title
250 def getStarPixCoord(self, exp, doMotionCorrection=True, useQfm=False):
251 target = exp.visitInfo.object
253 if self.useQfmForCentroids:
254 try:
255 result = self.qfmTask.run(exp)
256 pixCoord = result.brightestObjCentroid
257 expId = exp.info.id
258 logger.info(f'expId {expId} has centroid {pixCoord}')
259 except Exception:
260 return None
261 else:
262 pixCoord = getTargetCentroidFromWcs(exp, target, doMotionCorrection=doMotionCorrection)
263 return pixCoord
265 def makePng(self, dataId, saveFilename):
266 if self.exists(saveFilename) and not self.remakePngs: # should not be possible due to prerun
267 assert False, f"Almost overwrote {saveFilename} - how is this possible?"
269 if self.debug:
270 logger.info(f"Creating {saveFilename}")
272 self.fig.clear()
274 # must always keep exp unsmoothed for the centroiding via qfm
275 try:
276 exp = self.butler.get(self.dataProductToPlot, dataId)
277 except Exception:
278 # oh no, that should never happen, but it does! Let's just skip
279 logger.warning(f'Skipped {dataId}, because {self.dataProductToPlot} retrieval failed!')
280 return
281 toDisplay = exp
282 if self.smoothImages:
283 toDisplay = exp.clone()
284 toDisplay = self._smoothExp(toDisplay, 2)
286 try:
287 self.disp.mtv(toDisplay.image, title=self._titleFromExp(exp, dataId))
288 self.disp.scale('asinh', 'zscale')
289 except RuntimeError: # all-nan images slip through and don't display
290 self.disp.scale('linear', 0, 1)
291 self.disp.mtv(toDisplay.image, title=self._titleFromExp(exp, dataId))
292 self.disp.scale('asinh', 'zscale') # set back for next image
293 pass
295 if self.plotObjectCentroids:
296 try:
297 pixCoord = self.getStarPixCoord(exp)
298 if pixCoord:
299 self.disp.dot('x', *pixCoord, ctype='C1', size=50)
300 self.disp.dot('o', *pixCoord, ctype='C1', size=50)
301 else:
302 self.disp.dot('x', 2000, 2000, ctype='red', size=2000)
303 except Exception:
304 logger.warning(f"Failed to find OBJECT location for {dataId}")
306 deltaH = -0.05
307 deltaV = -0.05
308 plt.subplots_adjust(right=1+deltaH, left=0-deltaH, top=1+deltaV, bottom=0-deltaV)
309 self.fig.savefig(saveFilename)
310 logger.info(f'Saved png for {dataId} to {saveFilename}')
312 def pngsToMp4(self, indir, outfile, framerate, verbose=False):
313 """Create the movie with ffmpeg, from files."""
314 # NOTE: the order of ffmpeg arguments *REALLY MATTERS*.
315 # Reorder them at your own peril!
316 pathPattern = f'\"{os.path.join(indir, "*.png")}\"'
317 if verbose:
318 ffmpeg_verbose = 'info'
319 else:
320 ffmpeg_verbose = 'error'
321 cmd = ['ffmpeg',
322 '-v', ffmpeg_verbose,
323 '-f', 'image2',
324 '-y',
325 '-pattern_type glob',
326 '-framerate', f'{framerate}',
327 '-i', pathPattern,
328 '-vcodec', 'libx264',
329 '-b:v', '20000k',
330 '-profile:v', 'main',
331 '-pix_fmt', 'yuv420p',
332 '-threads', '10',
333 '-r', f'{framerate}',
334 os.path.join(outfile)]
336 subprocess.check_call(r' '.join(cmd), shell=True)
338 def tidyUp(self, tempDir):
339 shutil.rmtree(tempDir)
340 return
342 def _smoothExp(self, exp, smoothing, kernelSize=7):
343 """Use for DISPLAY ONLY!
345 Return a smoothed copy of the exposure
346 with the original mask plane in place."""
347 psf = measAlg.DoubleGaussianPsf(kernelSize, kernelSize, smoothing/(2*math.sqrt(2*math.log(2))))
348 newExp = exp.clone()
349 originalMask = exp.mask
351 kernel = psf.getKernel()
352 afwMath.convolve(newExp.maskedImage, newExp.maskedImage, kernel, afwMath.ConvolutionControl())
353 newExp.mask = originalMask
354 return newExp
357def animateDay(butler, dayObs, outputPath, dataProductToPlot='quickLookExp'):
358 outputFilename = f'{dayObs}.mp4'
360 onSkyIds = getLatissOnSkyDataIds(butler, startDate=dayObs, endDate=dayObs)
361 logger.info(f"Found {len(onSkyIds)} on sky ids for {dayObs}")
363 onSkyIds = [updateDataIdOrDataCord(dataId, detector=0) for dataId in onSkyIds]
365 animator = Animator(butler, onSkyIds, outputPath, outputFilename,
366 dataProductToPlot=dataProductToPlot,
367 remakePngs=False,
368 debug=False,
369 clobberVideoAndGif=True,
370 plotObjectCentroids=True,
371 useQfmForCentroids=True)
372 filename = animator.run()
373 return filename
376if __name__ == '__main__': 376 ↛ 377line 376 didn't jump to line 377, because the condition on line 376 was never true
377 outputPath = '/home/mfl/animatorOutput/main/'
378 butler = makeDefaultLatissButler()
380 day = 20211104
381 animateDay(butler, day, outputPath)