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

223 statements  

« prev     ^ index     » next       coverage.py v6.5.0, created at 2023-02-15 04:10 -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/>. 

21 

22import os 

23import subprocess 

24import shutil 

25import uuid 

26import math 

27 

28import matplotlib.pyplot as plt 

29 

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) 

40 

41from lsst.atmospec.utils import airMassFromRawMetadata 

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

43setupLogging() 

44 

45 

46class Animator(): 

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

48 for the data product specified.""" 

49 

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

59 

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

67 

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 

76 

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" 

81 

82 qfmTaskConfig = QuickFrameMeasurementTaskConfig() 

83 self.qfmTask = QuickFrameMeasurementTask(config=qfmTaskConfig) 

84 

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

90 

91 self.pngsToMakeDataIds = [] 

92 self.preRun() # sets the above list 

93 

94 @staticmethod 

95 def _strDataId(dataId): 

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

97 

98 Parameters 

99 ---------- 

100 dataId : `dict` 

101 The data id. 

102 

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

110 

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 

120 

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

122 """Convert dataId to filename. 

123 

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 

128 

129 dIdStr = self._strDataId(dataId) 

130 

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) 

137 

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

142 

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

148 

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

154 

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

161 

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

165 

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

169 

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

177 

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] 

182 

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) 

187 

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

195 

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 

200 

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) 

213 

214 # # create gif in temp dir 

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

216 # self.pngsToGif(pngFileList, outputGifFilename) 

217 

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) 

222 

223 # self.tidyUp(tempDir) 

224 # logger.info('Finished!') 

225 

226 # create gif in temp dir 

227 

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 

233 

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 

251 

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

253 target = exp.visitInfo.object 

254 

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 

266 

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

270 

271 if self.debug: 

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

273 

274 self.fig.clear() 

275 

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) 

287 

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 

296 

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

307 

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

313 

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

337 

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

339 

340 def tidyUp(self, tempDir): 

341 shutil.rmtree(tempDir) 

342 return 

343 

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

345 """Use for DISPLAY ONLY! 

346 

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 

352 

353 kernel = psf.getKernel() 

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

355 newExp.mask = originalMask 

356 return newExp 

357 

358 

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

360 outputFilename = f'{dayObs}.mp4' 

361 

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

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

364 

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

366 

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 

376 

377 

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

381 

382 day = 20211104 

383 animateDay(butler, day, outputPath)