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