def array2avi(A, fname, fps=25, codec='rawvideo', codec_params='', width=800,
makelog=False):
if A.ndim == 3:
fmt = 'gray'
elif A.ndim == 4:
if A.shape[3] == 3:
fmt = 'rgb8'
elif A.shape[3] == 4:
fmt = 'argb'
else:
raise Exception('Invalid input 4th dimension - should'
' be rgb or argb')
else:
raise Exception('Invalid number of input dimensions (should be'
' [frame,row,col] or [frame,row,col,(a)rgb]')
A = rescale_8bit(A)
nframes, h, w = A.shape[:3]
aspect = float(h) / w
width = 2 * int(width / 2.)
height = 2 * int((width / 2.) * aspect)
cmd = ['ffmpeg', '-y',
'-f', 'rawvideo',
'-pix_fmt', fmt,
'-r', '%d' % fps,
'-s', '%ix%i' % (w, h),
'-i', 'pipe:0',
'-vcodec', codec,
# '-b:v', '%ik' %kbitrate,
'-s', '%ix%i' % (width, height),
fname]
if makelog:
logfile = open(fname + '.log', 'w')
else:
logfile = open(devnull, 'w')
proc = sp.Popen(cmd, shell=False, stdin=sp.PIPE, stdout=logfile,
stderr=logfile)
wbh = Waitbar(title='Writing %s' % fname, showETA=True)
for ii, frame in enumerate(A):
proc.stdin.write(frame.tostring())
wbh.update((ii + 1.) / nframes)
proc.stdin.close()
proc.terminate()
logfile.close()
def save(self, fname, codec='rawvideo', fps=None, width=800,
makelog=False):
if not np.all(self.hasplayed):
raise Exception(
"All frames need to have been displayed at "
"least once before writing")
if fps is None:
fps = self.framerate
# make sure we start with the first frame
self.rewind()
# need to disconnect the first draw callback, since we'll be
# doing draws. otherwise, we'll end up starting the animation.
if self.animator._first_draw_id is not None:
self.animator._fig.canvas.mpl_disconnect(
self.animator._first_draw_id)
reconnect_first_draw = True
else:
reconnect_first_draw = False
# input pixel dimensions
w, h = self.fig.canvas.get_width_height()
nframes = self.data.shape[0]
# make sure that output width and height are divisible by 2
# (some codecs don't like odd dimensions)
figsize = self.fig.get_size_inches()
aspect = figsize[1] / figsize[0]
width = 2 * int(width / 2.)
height = 2 * int((width / 2.) * aspect)
# spawn an ffmpeg process - we're going to pipe the frames
# directly to ffmpeg as raw bytes
cmd = [
# global options
'ffmpeg', '-y',
'-loglevel', 'verbose',
# input options
'-f', 'rawvideo',
'-pix_fmt', 'bgr24',
'-r', '%d' % fps,
'-s', '%ix%i' % (w, h),
'-i', 'pipe:0',
# output options
'-vcodec', codec,
'-an',
'-pix_fmt', 'rgb24', # NB: need to reverse channel order!
'-s', '%ix%i' % (width, height),
fname
]
# print ' '.join(cmd)
if makelog:
logfile = open(fname + '.log', 'w')
else:
logfile = open(devnull, 'w')
proc = sp.Popen(cmd, shell=False, stdin=sp.PIPE,
stdout=logfile, stderr=logfile)
# Render each frame, save it to the stdin of the spawned process
wbh = Waitbar(title='Writing %s' % fname, showETA=True)
for ii, data in enumerate(self.animator.new_saved_frame_seq()):
self.animator._draw_next_frame(data, blit=True)
# self.animator._fig.savefig(proc.stdin,format='png')
proc.stdin.write(self.animator._fig.canvas.tostring_rgb())
wbh.update((ii + 1.) / nframes)
logfile.close()
if reconnect_first_draw:
drawid = self.animator._fig.canvas.mpl_connect(
'draw_event', self.animator._start)
self.animator._first_draw_id = drawid
