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