Coverage for python/lsst/summit/extras/animation.py: 12%

227 statements  

« prev     ^ index     » next       coverage.py v7.3.2, created at 2023-10-25 18:58 +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/>. 

21 

22import os 

23import subprocess 

24import shutil 

25import uuid 

26import math 

27import gc 

28 

29import matplotlib.pyplot as plt 

30 

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) 

41 

42from lsst.atmospec.utils import airMassFromRawMetadata 

43logger = logging.getLogger("lsst.summit.extras.animation") 

44setupLogging() 

45 

46 

47class Animator(): 

48 """Animate the list of dataIds in the order in which they are specified 

49 for the data product specified.""" 

50 

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): 

60 

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/") 

68 

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 

77 

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" 

82 

83 qfmTaskConfig = QuickFrameMeasurementTaskConfig() 

84 self.qfmTask = QuickFrameMeasurementTask(config=qfmTaskConfig) 

85 

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') 

91 

92 self.pngsToMakeDataIds = [] 

93 self.preRun() # sets the above list 

94 

95 @staticmethod 

96 def _strDataId(dataId): 

97 """Make a dataId into a string suitable for use as a filename. 

98 

99 Parameters 

100 ---------- 

101 dataId : `dict` 

102 The data id. 

103 

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}" 

111 

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 

121 

122 def dataIdToFilename(self, dataId, includeNumber=False, imNum=None): 

123 """Convert dataId to filename. 

124 

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 

129 

130 dIdStr = self._strDataId(dataId) 

131 

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) 

138 

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") 

143 

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}" 

149 

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") 

155 

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}") 

162 

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)] 

166 

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") 

170 

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") 

178 

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] 

183 

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) 

188 

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)) 

196 

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 

201 

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) 

214 

215 # # create gif in temp dir 

216 # outputGifFilename = os.path.join(tempDir, 'animation.gif') 

217 # self.pngsToGif(pngFileList, outputGifFilename) 

218 

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) 

223 

224 # self.tidyUp(tempDir) 

225 # logger.info('Finished!') 

226 

227 # create gif in temp dir 

228 

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 

234 

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 

252 

253 def getStarPixCoord(self, exp, doMotionCorrection=True, useQfm=False): 

254 target = exp.visitInfo.object 

255 

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 

267 

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?" 

271 

272 if self.debug: 

273 logger.info(f"Creating {saveFilename}") 

274 

275 self.fig.clear() 

276 

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) 

288 

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 

297 

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}") 

308 

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}') 

314 

315 del toDisplay 

316 del exp 

317 gc.collect() 

318 

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)] 

342 

343 subprocess.check_call(r' '.join(cmd), shell=True) 

344 

345 def tidyUp(self, tempDir): 

346 shutil.rmtree(tempDir) 

347 return 

348 

349 def _smoothExp(self, exp, smoothing, kernelSize=7): 

350 """Use for DISPLAY ONLY! 

351 

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 

357 

358 kernel = psf.getKernel() 

359 afwMath.convolve(newExp.maskedImage, newExp.maskedImage, kernel, afwMath.ConvolutionControl()) 

360 newExp.mask = originalMask 

361 return newExp 

362 

363 

364def animateDay(butler, dayObs, outputPath, dataProductToPlot='quickLookExp'): 

365 outputFilename = f'{dayObs}.mp4' 

366 

367 onSkyIds = getLatissOnSkyDataIds(butler, startDate=dayObs, endDate=dayObs) 

368 logger.info(f"Found {len(onSkyIds)} on sky ids for {dayObs}") 

369 

370 onSkyIds = [updateDataIdOrDataCord(dataId, detector=0) for dataId in onSkyIds] 

371 

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 

381 

382 

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() 

386 

387 day = 20211104 

388 animateDay(butler, day, outputPath)