This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| import argparse # need to install argparse: sudo apt-get install python-argparse | |
| import os, ast | |
| import numpy as np | |
| import motmot.FlyMovieFormat.FlyMovieFormat as FMF | |
| import pylab | |
| pylab.ion() | |
| import roslib | |
| roslib.load_manifest('rosbag') | |
| import rosbag | |
| #================================ function to load .bag files ====================================== | |
| def loadBag(inFileName): | |
| """ | |
| Script to import strokelitude bag file | |
| usage: outDict = loadBag(inFileName) | |
| PTW 02/04/12 | |
| """ | |
| if inFileName[-4:] != '.bag': | |
| inFileName = inFileName + '.bag' | |
| bag = rosbag.Bag(inFileName) | |
| timeSec = list() | |
| leftWingRad = list() | |
| rightWingRad = list() | |
| for topic, msg, t in bag.read_messages(): | |
| if topic == '/strokelitude' or topic == 'strokelitude': | |
| timeSec.append(msg.header.stamp.secs + msg.header.stamp.nsecs*float(1e-9)) | |
| leftWingRad.append(msg.left_wing_angle_radians) | |
| rightWingRad.append(msg.right_wing_angle_radians) | |
| timeSec = np.array(timeSec) | |
| leftWingRad=np.array(leftWingRad) | |
| rightWingRad = np.array(rightWingRad) | |
| outDict = {'timeSec':timeSec, 'leftWingRad':leftWingRad, 'rightWingRad':rightWingRad} | |
| return outDict | |
| def main(): | |
| """ it can be a bit tricky to enter the mask data in at the command line. Try copying this line and changing the values to what you need: | |
| python makeFmfStrokelitudeMovie.py "{'x':310.0, 'y':202.0, 'view_from_below':False, 'wingsplit':54.0, 'gamma':90.0}" | |
| """ | |
| #================================ Parse arguments ====================================== | |
| parser = argparse.ArgumentParser(description='Save a movie with strokelitude wingstroke amplitude tracking superimposed on original frames from .fmf file', epilog='PTW 2012') | |
| parser.add_argument('strokelitudeMaskData', type=str, help="dict of strokelitude mask for the .bag data file, e.g. {'x':310.0, 'y':202.0, 'view_from_below':False, 'wingsplit':54.0, 'gamma':90.0}") | |
| parser.add_argument('inFmfFileName', nargs='?', type=str, default=None, help='the .fmf file containing original movie frames') | |
| parser.add_argument('inBagFileName', nargs='?', type=str, default=None, help='the .bag file containing strokelitude tracking data') | |
| parser.add_argument('outFileName', nargs='?', type=str, default=None, help='the out file name') | |
| parser.add_argument('-s','--startTime', type=float, default=0, help='time in seconds from start of .fmf to start saved movie') | |
| parser.add_argument('-e','--endTime', type=float, default=0, help='time in seconds from start of .fmf to end saved movie') | |
| parser.add_argument('-r','--frameRateMultiplicationFactor', type=float, default=1, help='multiplicative factor to change frame rate from input to output movie') | |
| args = parser.parse_args() | |
| fileNames = os.listdir(os.getcwd()) | |
| fileNames.sort() | |
| if args.inFmfFileName is None: | |
| fmfFileNames = [f for f in fileNames if f[-3:].lower()=='fmf'] | |
| if len(fmfFileNames) > 0: | |
| print 'since no input .fmf file specified, using', fmfFileNames[0] | |
| args.inFmfFileName = fmfFileNames[0] | |
| else: | |
| parser.print_help() | |
| parser.error('need to specify .fmf file name') | |
| if args.inBagFileName is None: | |
| bagFileNames = [f for f in fileNames if f[-3:].lower()=='bag'] | |
| if len(bagFileNames) > 0: | |
| print 'since no input .bag file specified, using', bagFileNames[0] | |
| args.inBagFileName = bagFileNames[0] | |
| else: | |
| parser.print_help() | |
| parser.error('need to specify .bag file name') | |
| if args.inFmfFileName[-4:].lower() != '.fmf': | |
| args.inFmfFileName = args.inFmfFileName+'.fmf' | |
| if args.inBagFileName[-4:].lower() != '.bag': | |
| args.inBagFileName = args.inBagFileName+'.bag' | |
| if args.outFileName is None: | |
| args.outFileName = 'out'+args.inFmfFileName[:-4]+args.inBagFileName[:-4]+'.mpeg' | |
| print 'no output name specified, using', args.outFileName | |
| if args.outFileName[-5:].lower() != '.mpeg': | |
| args.outFileName = args.outFileName+'.mpeg' | |
| print 'in fmf file name is', args.inFmfFileName | |
| print 'in bag file name is', args.inBagFileName | |
| print 'out file name is', args.outFileName | |
| #================================ strokelitude conversion stuff ====================================== | |
| strokelitudeMaskData = ast.literal_eval(args.strokelitudeMaskData) | |
| D2R = np.pi/180.0 | |
| mult_sign={True:{'left':1, | |
| 'right':-1}, | |
| False:{'left':-1, | |
| 'right':1}} | |
| gammaRad = strokelitudeMaskData['gamma']*D2R | |
| rotation = np.array([[ np.cos( gammaRad ), -np.sin(gammaRad)],[ np.sin( gammaRad ), np.cos(gammaRad)]]) | |
| translation = np.array([[strokelitudeMaskData['x']],[strokelitudeMaskData['y']]]) | |
| def get_wingsplit_translation(side): | |
| sign = mult_sign[strokelitudeMaskData['view_from_below']][side] | |
| return np.array([[0.0],[sign*strokelitudeMaskData['wingsplit']]]) | |
| def get_span_lineseg(side,theta_user): | |
| """draw line on side at angle theta_user (in radians)""" | |
| linesegs = [] | |
| sign = mult_sign[strokelitudeMaskData['view_from_below']][side] | |
| theta = 90*D2R - theta_user | |
| theta *= sign | |
| verts = np.array( [[ 0, 1000.0*np.cos(theta)], | |
| [ 0, 1000.0*np.sin(theta)]] ) | |
| verts = verts + get_wingsplit_translation(side) | |
| verts = np.dot(rotation, verts) + translation | |
| linesegs.append( verts.T.ravel() ) | |
| return linesegs | |
| # ======= read bag strokelitude data ============== | |
| allBagData = loadBag(args.inBagFileName) | |
| bagTime = allBagData['timeSec'].copy() | |
| leftWingRad = allBagData['leftWingRad'].copy() | |
| rightWingRad = allBagData['rightWingRad'].copy() | |
| os.system('mkdir TEMPDIRFORMOVIEFRAMES') | |
| outDirName = os.path.join(os.getcwd(),'TEMPDIRFORMOVIEFRAMES') | |
| # ======= open fmf file=========================== | |
| infmffile = FMF.FlyMovie(args.inFmfFileName) | |
| numFlyFrames = infmffile.get_n_frames() | |
| flyTime = infmffile.get_all_timestamps() | |
| flyFrameRate = 1.0/np.mean(np.diff(flyTime)) | |
| if flyFrameRate*args.frameRateMultiplicationFactor < 15: | |
| print 'ERROR: cannot have output frame rate below 15 fps, your frameRateMultiplicationFactor of', args.frameRateMultiplicationFactor, 'combined with the input frame rate of', flyFrameRate, 'results in a desired frame rate of', flyFrameRate*args.frameRateMultiplicationFactor | |
| return | |
| inFMFHeight = infmffile.get_height() | |
| inFMFWidth = infmffile.get_width() | |
| flyFrameStartTime = flyTime[0] + args.startTime | |
| flyFrameStopTime = flyTime[0] + args.endTime | |
| flyStartInd = np.argmin(abs(flyTime - flyFrameStartTime)) | |
| flyStopInd = np.argmin(abs(flyTime - flyFrameStopTime)) | |
| bagStopInd = np.argmin(abs(bagTime - flyFrameStopTime)) | |
| fig = pylab.figure(figsize=(8,6)) | |
| fig.set_facecolor('w') | |
| LEFTCOLOR = 'b' | |
| RIGHTCOLOR = [0,1,0] | |
| #fig.text(.05,.96,'10x normal speed, diffuser above polarizer') | |
| flyAxes = fig.add_axes([0.15, 0.28, 0.7, 0.7]) | |
| flyAxes.axis('image') | |
| wingsTimePlotAxes = fig.add_axes([0.15, 0.1, 0.7, 0.15]) | |
| wingsTimePlotAxes.plot(bagTime-flyFrameStartTime,leftWingRad*180/np.pi,color=LEFTCOLOR,linewidth=.5,zorder=10) | |
| wingsTimePlotAxes.plot(bagTime-flyFrameStartTime,rightWingRad*180/np.pi,color=RIGHTCOLOR,linewidth=.5,zorder=10) | |
| wingsTimePlotAxes.set_ylim((0,90)) | |
| wingsTimePlotAxes.spines['bottom'].set_position(('outward',8)) | |
| wingsTimePlotAxes.spines['bottom'].set_position(('outward',8)) | |
| wingsTimePlotAxes.xaxis.set_ticks_position('bottom') | |
| wingsTimePlotAxes.set_xlim((0,flyFrameStopTime-flyFrameStartTime)) | |
| wingsTimePlotAxes.spines['left'].set_position(('outward',8)) | |
| wingsTimePlotAxes.spines['top'].set_visible(False) | |
| wingsTimePlotAxes.spines['right'].set_visible(False) | |
| wingsTimePlotAxes.yaxis.set_ticks_position('left') | |
| wingsTimePlotAxes.set_yticks([0,45,90]) | |
| wingsTimePlotAxes.set_yticklabels(['0$^\circ$','45$^\circ$','90$^\circ$']) | |
| wingsTimePlotAxes.set_ylabel('Wingstroke \namplitude') | |
| wingsTimePlotAxes.set_xlabel('Time from start (s)') | |
| wingsTimePlotAxes.text(wingsTimePlotAxes.get_xlim()[1],leftWingRad[bagStopInd]*180/np.pi,'Left',color=LEFTCOLOR) | |
| wingsTimePlotAxes.text(wingsTimePlotAxes.get_xlim()[1],rightWingRad[bagStopInd]*180/np.pi,'Right',color=RIGHTCOLOR) | |
| flyAxes.cla() | |
| for flyIndThisFrameInd, flyIndThisFrame in enumerate(range(flyStartInd, flyStopInd)): | |
| print 'processing frame ', flyIndThisFrameInd, ' of ', flyStopInd - flyStartInd | |
| flyTimeThisFrame = flyTime[flyIndThisFrame] | |
| bagIndThisFrame = np.argmin(abs(bagTime - flyTimeThisFrame)) | |
| bagTimeThisFrame = bagTime[bagIndThisFrame] | |
| flyFrame,flyTimeThisFrameFromFMF = infmffile.get_frame(flyIndThisFrame) | |
| if flyTimeThisFrameFromFMF != flyTimeThisFrame: | |
| print 'different fly times this frame??' | |
| flyAxes.imshow(flyFrame,cmap='gray',origin='lower') | |
| leftLineSegs = get_span_lineseg('left',leftWingRad[bagIndThisFrame]) | |
| rightLineSegs = get_span_lineseg('right',rightWingRad[bagIndThisFrame]) | |
| flyAxes.plot([leftLineSegs[0][0],leftLineSegs[0][2]],[leftLineSegs[0][1],leftLineSegs[0][3]],color=LEFTCOLOR,linewidth=2) | |
| flyAxes.plot([rightLineSegs[0][0],rightLineSegs[0][2]],[rightLineSegs[0][1],rightLineSegs[0][3]],color=RIGHTCOLOR,linewidth=2) | |
| flyAxes.axis('off') | |
| flyAxes.set_xlim((0,flyFrame.shape[1])) | |
| flyAxes.set_ylim((0,flyFrame.shape[0])) | |
| movingWingsLine = wingsTimePlotAxes.axvline(bagTimeThisFrame-flyFrameStartTime,color='k') | |
| wingsTimePlotAxes.set_xlim((0,flyFrameStopTime-flyFrameStartTime)) | |
| fig.savefig(os.path.join(outDirName,'frame')+("%05d" %(flyIndThisFrameInd+1))+'.png') | |
| if flyIndThisFrame < flyStopInd - 1: | |
| flyAxes.cla() | |
| movingWingsLine.remove() | |
| ffmpegCommandString = 'ffmpeg -b 8000000 -r '+str(flyFrameRate*args.frameRateMultiplicationFactor)+' -i '+os.path.join(outDirName,'frame%05d.png ')+args.outFileName | |
| #print ffmpegCommandString | |
| os.system(ffmpegCommandString) | |
| #=================== clean up temp folders and directory ========================== | |
| os.system('rm '+os.path.join(outDirName,'frame*.png ')) | |
| os.system('rmdir '+outDirName) | |
| return | |
| if __name__ == "__main__": | |
| main() |
No comments:
Post a Comment