Python matplotlib.animation.writers() Examples

def save_frames(images, filename):
    num_sequences, n_steps, w, h = images.shape

    fig = plt.figure()
    im = plt.imshow(combine_multiple_img(images[:, 0]), cmap=plt.cm.get_cmap('Greys'), interpolation='none')
    plt.axis('image')

    def updatefig(*args):
        im.set_array(combine_multiple_img(images[:, args[0]]))
        return im,

    ani = animation.FuncAnimation(fig, updatefig, interval=500, frames=n_steps)

    # Either avconv or ffmpeg need to be installed in the system to produce the videos!
    try:
        writer = animation.writers['avconv']
    except KeyError:
        writer = animation.writers['ffmpeg']
    writer = writer(fps=3)
    ani.save(filename, writer=writer)
    plt.close(fig) 

def make_video(xy,filename):
  FFMpegWriter = manimation.writers['ffmpeg']
  metadata = dict(title='Movie Test', artist='Matplotlib',
                  comment='Movie support!')
  writer = FFMpegWriter(fps=15, metadata=metadata)
  mydpi=100;
  #fig = plt.figure(figsize=(128/mydpi,128/mydpi))
  fig = plt.figure(figsize=(32/mydpi,32/mydpi))
  plt.xlim(-200, 200)
  plt.ylim(-200, 200)
  fig_num=len(xy);
  #color=['ro','bo','go','ko','yo','mo','co'];
  color=['r','b','g','k','y','m','c'];
  with writer.saving(fig, filename, len(xy)):
    for i in range(len(xy)):
      for j in range(len(xy[0])):
        #plt.plot(xy[i,j,1],xy[i,j,0],color[j%len(color)]);
        plt.scatter(xy[i,j,1],xy[i,j,0],c=color[j%len(color)],s=0.5);
      writer.grab_frame(); 

def make_video(xy,filename):
  os.system("rm -rf pics/*");
  FFMpegWriter = manimation.writers['ffmpeg']
  metadata = dict(title='Movie Test', artist='Matplotlib',
                  comment='Movie support!')
  writer = FFMpegWriter(fps=15, metadata=metadata)
  fig = plt.figure()
  plt.xlim(-200, 200)
  plt.ylim(-200, 200)
  fig_num=len(xy);
  color=['ro','bo','go','ko','yo','mo','co'];
  with writer.saving(fig, filename, len(xy)):
    for i in range(len(xy)):
      for j in range(len(xy[0])):
        plt.plot(xy[i,j,1],xy[i,j,0],color[j%len(color)]);
      writer.grab_frame(); 

def coalescence_video(self, file_str):
        """
        Generate a video over the marginal window showing the coalescence map
        and expected arrival times overlain on the station traces.

        Parameters
        ----------
        file_str : str
            String {run_name}_{event_name}

        """

        # Find index of start and end of marginal window
        idx0 = np.where(self.times == self.event_mw_data["DT"].iloc[0])[0][0]
        idx1 = np.where(self.times == self.event_mw_data["DT"].iloc[-1])[0][0]

        Writer = animation.writers["ffmpeg"]
        writer = Writer(fps=4, metadata=dict(artist="Ulvetanna"), bitrate=1800)

        fig = self._coalescence_frame(idx0)
        ani = animation.FuncAnimation(fig, self._video_update,
                                      frames=np.linspace(idx0+1, idx1, 200),
                                      blit=False, repeat=False)

        subdir = "videos"
        util.make_directories(self.run_path, subdir=subdir)
        out_str = self.run_path / subdir / file_str
        ani.save("{}_CoalescenceVideo.mp4".format(out_str),
                 writer=writer) 

def savemp4_per_desc(self, fname):
        if not os.path.isdir(self.loss_name):
            os.makedirs(self.loss_name)
        img = self.img1
        for d_idx in range(len(self.names_list)):
            fig, ax = plt.subplots()
            fig.set_tight_layout(True)
            ax.imshow(255 - img.squeeze())
            lines_dict = {}
            N = self.names_list[d_idx];
            C = self.colors[d_idx];
            lafs_dict = {"1": self.out_lafs1[N], "2": self.out_lafs2[N]}
            num_frames = len(lafs_dict["1"])
            lines_dict['1'], ax =  visualize_LAFs_on_fig(ax, lafs_dict["1"][0], 'r')
            lines_dict['2'], ax =  visualize_LAFs_on_fig(ax, lafs_dict["2"][0], 'b')
            ax.legend(['img1', 'img2'])
            def visualize_LAFs_on_fig_mod(laf_dict, line_dict, dname, i):
                work_LAFs = convertLAFs_to_A23format(laf_dict[dname][i])
                for jj in range(len(work_LAFs)):
                    ell = LAF2pts(work_LAFs[jj,:,:])
                    line_dict[dname][str(jj)].set_data(ell[:,0], ell[:,1])
                return line_dict[dname],ax
            def update_LAF(i):
                lines_dict["1"],ax =  visualize_LAFs_on_fig_mod(lafs_dict, lines_dict, "1", i)
                lines_dict["2"],ax =  visualize_LAFs_on_fig_mod(lafs_dict, lines_dict, "2", i)
                return lines_dict["1"], ax 
            anim = FuncAnimation(fig, update_LAF, frames=np.arange(0, num_frames), interval=75)
            Writer = animation.writers['ffmpeg']
            writer = Writer(fps=24, metadata=dict(artist='Me'), bitrate=900)
            anim.save(os.path.join(self.loss_name, N + "_" +  fname), dpi=72, writer=writer)
        return 

def savemp4_per_img(self, img, lafs_dict, fname):
        fig, ax = plt.subplots()
        fig.set_tight_layout(True)
        ax.imshow(255 - img.squeeze())
        lines_dict = {}
        for d_idx in range(len(self.names_list)):
            N = self.names_list[d_idx];
            C = self.colors[d_idx];
            num_frames = len(lafs_dict[N])
            lines_dict[N], ax =  visualize_LAFs_on_fig(ax, lafs_dict[N][0], C)
        ax.legend(self.names_list)
        def visualize_LAFs_on_fig_mod(laf_dict, line_dict, dname, i):
            work_LAFs = convertLAFs_to_A23format(laf_dict[dname][i])
            ells = []
            for jj in range(len(work_LAFs)):
                ell = LAF2pts(work_LAFs[jj,:,:])
                line_dict[dname][str(i)].set_data(ell[:,0], ell[:,1])
            return line_dict[dname],ax
        def update_LAF(i):
            for d_idx in range(len(self.names_list)):
                N = self.names_list[d_idx];
                C = self.colors[d_idx];
                lines_dict[N],ax =  visualize_LAFs_on_fig_mod(lafs_dict, lines_dict, N, i)
            return lines_dict[N], ax 
        ax.legend(self.names_list)
        anim = FuncAnimation(fig, update_LAF, frames=np.arange(0, num_frames), interval=75)
        import matplotlib.animation as animation
        Writer = animation.writers['ffmpeg']
        writer = Writer(fps=24, metadata=dict(artist='Me'), bitrate=1800)
        anim.save(fname, dpi=96, writer=writer)#'imagemagick')        
        return 

def generate_writer():
    FFMpegWriter = animation.writers['ffmpeg']
    writer = FFMpegWriter(fps=parameters.frames_per_second, metadata=parameters.metadata)
    fig = pyplot.figure()
    fig.subplots_adjust(left=0, bottom=0, right=1, top=1, wspace=None, hspace=None)
    pyplot.xlim(0, parameters.width)
    pyplot.ylim(0, parameters.height)
    axis = pyplot.gca()
    axis.set_axis_bgcolor('black')
    axis.axes.get_xaxis().set_visible(False)
    axis.axes.get_yaxis().set_visible(False)
    rcParams['font.size'] = 12
    rcParams['text.color'] = 'white'
    return fig, writer 

def save_traj(images, image_goal, gif_path, task):
    # save trajectory as gif file
    fig, aa = plt.subplots(1, 2)
    m1 = aa[0].matshow(
        images[0], cmap=plt.cm.gray, vmin=0., vmax=1.)
    aa[0].set_title('Time step 0')
    aa[0].set_yticklabels([])
    aa[0].set_xticklabels([])
    m2 = aa[1].matshow(
        image_goal, cmap=plt.cm.gray, vmin=0., vmax=1.)
    aa[1].set_title('goal')
    aa[1].set_yticklabels([])
    aa[1].set_xticklabels([])
    fig.tight_layout()

    def updatemat2(t):
        m1.set_data(images[t])
        aa[0].set_title('Time step ' + str(t))
        m2.set_data(image_goal)
        return m1, m2

    frames = len(images)
    if task == 'plane':
        fps = 2
    else:
        fps = 20

    anim = FuncAnimation(
        fig, updatemat2, frames=frames, interval=200, blit=True, repeat=True)
    Writer = writers['imagemagick']  # animation.writers.avail
    writer = Writer(fps=fps, metadata=dict(artist='Me'), bitrate=1800)

    anim.save(gif_path, writer=writer) 

def save_video(self):
        FFMpegWriter = manimation.writers['ffmpeg']
        options = dict(title='Movie', artist='Matplotlib')
        writer = FFMpegWriter(fps=5, options=options, bitrate=3000)

        with writer.saving(self.f, "movie.mp4", 100):
            self.curr_pos = 0
            self.update_graphs()
            self.f.canvas.draw()
            plt.draw()
            for i in range(10):
                writer.grab_frame()

            for i in range(self.num_iters):
                self.curr_pos = i
                self.update_graphs()
                self.f.canvas.draw()
                plt.draw()
                writer.grab_frame()

            self.curr_pos = self.num_iters
            self.update_graphs()
            self.f.canvas.draw()
            plt.draw()
            for i in range(20):
                writer.grab_frame() 

def save_video(self):
        FFMpegWriter = manimation.writers['ffmpeg']
        options = dict(title='Movie', artist='Matplotlib')
        writer = FFMpegWriter(fps=5, options=options, bitrate=3000)

        with writer.saving(self.f, "movie.mp4", 100):
            self.curr_pos = 0
            self.update_graphs()
            self.f.canvas.draw()
            plt.draw()
            for i in range(10):
                writer.grab_frame()

            for i in range(self.num_iters):
                self.curr_pos = i
                self.update_graphs()
                self.f.canvas.draw()
                plt.draw()
                writer.grab_frame()

            self.curr_pos = self.num_iters
            self.update_graphs()
            self.f.canvas.draw()
            plt.draw()
            for i in range(20):
                writer.grab_frame() 

def movie(im_List, out='movie.mp4', fps=10, dpi=120):
    import matplotlib
    matplotlib.use('agg')
    import matplotlib.pyplot as plt
    import matplotlib.animation as animation

    fig = plt.figure()
    frame = im_List[0].imvec #read_auto(filelist[len(filelist)/2])
    fov = im_List[0].psize*im_List[0].xdim
    extent = fov * np.array((1,-1,-1,1)) / 2.
    maxi = np.max(frame)
    im = plt.imshow( np.reshape(frame,[im_List[0].xdim, im_List[0].xdim]) , cmap='hot', extent=extent) #inferno
    plt.colorbar()
    im.set_clim([0,maxi])
    fig.set_size_inches([5,5])
    plt.tight_layout()

    def update_img(n):
        sys.stdout.write('\rprocessing image %i of %i ...' % (n,len(im_List)) )
        sys.stdout.flush()
        im.set_data(np.reshape(im_List[n].imvec, [im_List[n].xdim, im_List[n].xdim]) )
        return im

    ani = animation.FuncAnimation(fig,update_img,len(im_List),interval=1e3/fps)
    writer = animation.writers['ffmpeg'](fps=max(20, fps), bitrate=1e6)
    ani.save(out,writer=writer,dpi=dpi) 

def save_true_generated_frames(true, generated, filename):
    num_sequences, n_steps, w, h = true.shape

    # Background is 0, foreground as 1
    true = np.copy(true[:16])
    true[true > 0.1] = 1

    # Set foreground be near 0.5
    generated = generated * .5

    # Background is 1, foreground is near 0.5
    generated = 1 - generated[:16, :n_steps]

    # Subtract true from generated so background is 1, true foreground is 0,
    # and generated foreground is around 0.5
    images = generated - true
    # images[images > 0.5] = 1.

    fig = plt.figure()
    im = plt.imshow(combine_multiple_img(images[:, 0]), cmap=plt.cm.get_cmap('gist_heat'),
                    interpolation='none', vmin=0, vmax=1)
    plt.axis('image')

    def updatefig(*args):
        im.set_array(combine_multiple_img(images[:, args[0]]))
        return im,

    ani = animation.FuncAnimation(fig, updatefig, interval=500, frames=n_steps)

    try:
        writer = animation.writers['avconv']
    except KeyError:
        writer = animation.writers['ffmpeg']
    writer = writer(fps=3)
    ani.save(filename, writer=writer)
    plt.close(fig) 

def MakeImitationVideo(
    outdir, vidname, query_im_strs, knn_im_strs, height=640, width=360):
  """Creates a KNN imitation video.

  For each frame in vid0, pair with the frame at index in knn_indices in
  vids1. Write video to disk.

  Args:
    outdir: String, directory to write videos.
    vidname: String, name of video.
    query_im_strs: Numpy array holding query image strings.
    knn_im_strs: Numpy array holding knn image strings.
    height: Int, height of raw images.
    width: Int, width of raw images.
  """
  if not tf.gfile.Exists(outdir):
    tf.gfile.MakeDirs(outdir)
  vid_path = os.path.join(outdir, vidname)
  combined = zip(query_im_strs, knn_im_strs)

  # Create and write the video.
  fig = plt.figure()
  ax = fig.add_subplot(111)
  ax.set_aspect('equal')
  ax.get_xaxis().set_visible(False)
  ax.get_yaxis().set_visible(False)
  im = ax.imshow(
      np.zeros((height, width*2, 3)), cmap='gray', interpolation='nearest')
  im.set_clim([0, 1])
  plt.tight_layout(pad=0, w_pad=0, h_pad=0)
  # pylint: disable=invalid-name
  def update_img(pair):
    """Decode pairs of image strings, update a video."""
    im_i, im_j = pair
    nparr_i = np.fromstring(str(im_i), np.uint8)
    img_np_i = cv2.imdecode(nparr_i, 1)
    img_np_i = img_np_i[..., [2, 1, 0]]
    nparr_j = np.fromstring(str(im_j), np.uint8)
    img_np_j = cv2.imdecode(nparr_j, 1)
    img_np_j = img_np_j[..., [2, 1, 0]]

    # Optionally reshape the images to be same size.
    frame = np.concatenate([img_np_i, img_np_j], axis=1)
    im.set_data(frame)
    return im
  ani = animation.FuncAnimation(fig, update_img, combined, interval=15)
  writer = animation.writers['ffmpeg'](fps=15)
  dpi = 100
  tf.logging.info('Writing video to:\n %s \n' % vid_path)
  ani.save('%s.mp4' % vid_path, writer=writer, dpi=dpi) 

def MakeImitationVideo(
    outdir, vidname, query_im_strs, knn_im_strs, height=640, width=360):
  """Creates a KNN imitation video.

  For each frame in vid0, pair with the frame at index in knn_indices in
  vids1. Write video to disk.

  Args:
    outdir: String, directory to write videos.
    vidname: String, name of video.
    query_im_strs: Numpy array holding query image strings.
    knn_im_strs: Numpy array holding knn image strings.
    height: Int, height of raw images.
    width: Int, width of raw images.
  """
  if not tf.gfile.Exists(outdir):
    tf.gfile.MakeDirs(outdir)
  vid_path = os.path.join(outdir, vidname)
  combined = zip(query_im_strs, knn_im_strs)

  # Create and write the video.
  fig = plt.figure()
  ax = fig.add_subplot(111)
  ax.set_aspect('equal')
  ax.get_xaxis().set_visible(False)
  ax.get_yaxis().set_visible(False)
  im = ax.imshow(
      np.zeros((height, width*2, 3)), cmap='gray', interpolation='nearest')
  im.set_clim([0, 1])
  plt.tight_layout(pad=0, w_pad=0, h_pad=0)
  # pylint: disable=invalid-name
  def update_img(pair):
    """Decode pairs of image strings, update a video."""
    im_i, im_j = pair
    nparr_i = np.fromstring(str(im_i), np.uint8)
    img_np_i = cv2.imdecode(nparr_i, 1)
    img_np_i = img_np_i[..., [2, 1, 0]]
    nparr_j = np.fromstring(str(im_j), np.uint8)
    img_np_j = cv2.imdecode(nparr_j, 1)
    img_np_j = img_np_j[..., [2, 1, 0]]

    # Optionally reshape the images to be same size.
    frame = np.concatenate([img_np_i, img_np_j], axis=1)
    im.set_data(frame)
    return im
  ani = animation.FuncAnimation(fig, update_img, combined, interval=15)
  writer = animation.writers['ffmpeg'](fps=15)
  dpi = 100
  tf.logging.info('Writing video to:\n %s \n' % vid_path)
  ani.save('%s.mp4' % vid_path, writer=writer, dpi=dpi) 

def MakeImitationVideo(
    outdir, vidname, query_im_strs, knn_im_strs, height=640, width=360):
  """Creates a KNN imitation video.

  For each frame in vid0, pair with the frame at index in knn_indices in
  vids1. Write video to disk.

  Args:
    outdir: String, directory to write videos.
    vidname: String, name of video.
    query_im_strs: Numpy array holding query image strings.
    knn_im_strs: Numpy array holding knn image strings.
    height: Int, height of raw images.
    width: Int, width of raw images.
  """
  if not tf.gfile.Exists(outdir):
    tf.gfile.MakeDirs(outdir)
  vid_path = os.path.join(outdir, vidname)
  combined = zip(query_im_strs, knn_im_strs)

  # Create and write the video.
  fig = plt.figure()
  ax = fig.add_subplot(111)
  ax.set_aspect('equal')
  ax.get_xaxis().set_visible(False)
  ax.get_yaxis().set_visible(False)
  im = ax.imshow(
      np.zeros((height, width*2, 3)), cmap='gray', interpolation='nearest')
  im.set_clim([0, 1])
  plt.tight_layout(pad=0, w_pad=0, h_pad=0)
  # pylint: disable=invalid-name
  def update_img(pair):
    """Decode pairs of image strings, update a video."""
    im_i, im_j = pair
    nparr_i = np.fromstring(str(im_i), np.uint8)
    img_np_i = cv2.imdecode(nparr_i, 1)
    img_np_i = img_np_i[..., [2, 1, 0]]
    nparr_j = np.fromstring(str(im_j), np.uint8)
    img_np_j = cv2.imdecode(nparr_j, 1)
    img_np_j = img_np_j[..., [2, 1, 0]]

    # Optionally reshape the images to be same size.
    frame = np.concatenate([img_np_i, img_np_j], axis=1)
    im.set_data(frame)
    return im
  ani = animation.FuncAnimation(fig, update_img, combined, interval=15)
  writer = animation.writers['ffmpeg'](fps=15)
  dpi = 100
  tf.logging.info('Writing video to:\n %s \n' % vid_path)
  ani.save('%s.mp4' % vid_path, writer=writer, dpi=dpi) 

def animate(path_to_file, H, problem=None, func_iter=None, plot_min=None, plot_max=None):
    if H.ndim != 3 or H.shape[2] != 2:
        print("Can only animate a two dimensional set of arrays.")
        return

    fig = plt.figure()
    ax = plt.gca()

    # plot the pareto front if it is known for the problem
    if problem is not None:
        pf = problem.pareto_front()
        plt.scatter(pf[:, 0], pf[:, 1], label='Pareto Front', s=60, facecolors='none', edgecolors='r')

    # plot the initial population
    _F = H[0, :, :]
    scat = plt.scatter(_F[:, 0], _F[:, 1])
    plt.title("0")

    if func_iter is not None:
        func_iter(ax, H[0])

    # the update method
    def update(n):
        _F = H[n, :, :]
        scat.set_offsets(_F)

        # get the bounds for plotting and add padding
        min = np.min(_F, axis=0) - 0.1
        max = np.max(_F, axis=0) + 0.1

        # set the scatter object with padding
        ax.set_xlim(min[0], max[0])
        ax.set_ylim(min[1], max[1])

        if func_iter is not None:
            func_iter(ax, H[n])

        plt.title(n)

    # create the animation
    ani = animation.FuncAnimation(fig, update, frames=H.shape[0])

    # write the file
    Writer = animation.writers['ffmpeg']
    writer = Writer(fps=6, bitrate=1800)
    ani.save(path_to_file, writer=writer)

    print("Saving: ", path_to_file) 

def MakeImitationVideo(
    outdir, vidname, query_im_strs, knn_im_strs, height=640, width=360):
  """Creates a KNN imitation video.

  For each frame in vid0, pair with the frame at index in knn_indices in
  vids1. Write video to disk.

  Args:
    outdir: String, directory to write videos.
    vidname: String, name of video.
    query_im_strs: Numpy array holding query image strings.
    knn_im_strs: Numpy array holding knn image strings.
    height: Int, height of raw images.
    width: Int, width of raw images.
  """
  if not tf.gfile.Exists(outdir):
    tf.gfile.MakeDirs(outdir)
  vid_path = os.path.join(outdir, vidname)
  combined = zip(query_im_strs, knn_im_strs)

  # Create and write the video.
  fig = plt.figure()
  ax = fig.add_subplot(111)
  ax.set_aspect('equal')
  ax.get_xaxis().set_visible(False)
  ax.get_yaxis().set_visible(False)
  im = ax.imshow(
      np.zeros((height, width*2, 3)), cmap='gray', interpolation='nearest')
  im.set_clim([0, 1])
  plt.tight_layout(pad=0, w_pad=0, h_pad=0)
  # pylint: disable=invalid-name
  def update_img(pair):
    """Decode pairs of image strings, update a video."""
    im_i, im_j = pair
    nparr_i = np.fromstring(str(im_i), np.uint8)
    img_np_i = cv2.imdecode(nparr_i, 1)
    img_np_i = img_np_i[..., [2, 1, 0]]
    nparr_j = np.fromstring(str(im_j), np.uint8)
    img_np_j = cv2.imdecode(nparr_j, 1)
    img_np_j = img_np_j[..., [2, 1, 0]]

    # Optionally reshape the images to be same size.
    frame = np.concatenate([img_np_i, img_np_j], axis=1)
    im.set_data(frame)
    return im
  ani = animation.FuncAnimation(fig, update_img, combined, interval=15)
  writer = animation.writers['ffmpeg'](fps=15)
  dpi = 100
  tf.logging.info('Writing video to:\n %s \n' % vid_path)
  ani.save('%s.mp4' % vid_path, writer=writer, dpi=dpi) 

def MakeImitationVideo(
    outdir, vidname, query_im_strs, knn_im_strs, height=640, width=360):
  """Creates a KNN imitation video.

  For each frame in vid0, pair with the frame at index in knn_indices in
  vids1. Write video to disk.

  Args:
    outdir: String, directory to write videos.
    vidname: String, name of video.
    query_im_strs: Numpy array holding query image strings.
    knn_im_strs: Numpy array holding knn image strings.
    height: Int, height of raw images.
    width: Int, width of raw images.
  """
  if not tf.gfile.Exists(outdir):
    tf.gfile.MakeDirs(outdir)
  vid_path = os.path.join(outdir, vidname)
  combined = zip(query_im_strs, knn_im_strs)

  # Create and write the video.
  fig = plt.figure()
  ax = fig.add_subplot(111)
  ax.set_aspect('equal')
  ax.get_xaxis().set_visible(False)
  ax.get_yaxis().set_visible(False)
  im = ax.imshow(
      np.zeros((height, width*2, 3)), cmap='gray', interpolation='nearest')
  im.set_clim([0, 1])
  plt.tight_layout(pad=0, w_pad=0, h_pad=0)
  # pylint: disable=invalid-name
  def update_img(pair):
    """Decode pairs of image strings, update a video."""
    im_i, im_j = pair
    nparr_i = np.fromstring(str(im_i), np.uint8)
    img_np_i = cv2.imdecode(nparr_i, 1)
    img_np_i = img_np_i[..., [2, 1, 0]]
    nparr_j = np.fromstring(str(im_j), np.uint8)
    img_np_j = cv2.imdecode(nparr_j, 1)
    img_np_j = img_np_j[..., [2, 1, 0]]

    # Optionally reshape the images to be same size.
    frame = np.concatenate([img_np_i, img_np_j], axis=1)
    im.set_data(frame)
    return im
  ani = animation.FuncAnimation(fig, update_img, combined, interval=15)
  writer = animation.writers['ffmpeg'](fps=15)
  dpi = 100
  tf.logging.info('Writing video to:\n %s \n' % vid_path)
  ani.save('%s.mp4' % vid_path, writer=writer, dpi=dpi) 

def MakeImitationVideo(
    outdir, vidname, query_im_strs, knn_im_strs, height=640, width=360):
  """Creates a KNN imitation video.

  For each frame in vid0, pair with the frame at index in knn_indices in
  vids1. Write video to disk.

  Args:
    outdir: String, directory to write videos.
    vidname: String, name of video.
    query_im_strs: Numpy array holding query image strings.
    knn_im_strs: Numpy array holding knn image strings.
    height: Int, height of raw images.
    width: Int, width of raw images.
  """
  if not tf.gfile.Exists(outdir):
    tf.gfile.MakeDirs(outdir)
  vid_path = os.path.join(outdir, vidname)
  combined = zip(query_im_strs, knn_im_strs)

  # Create and write the video.
  fig = plt.figure()
  ax = fig.add_subplot(111)
  ax.set_aspect('equal')
  ax.get_xaxis().set_visible(False)
  ax.get_yaxis().set_visible(False)
  im = ax.imshow(
      np.zeros((height, width*2, 3)), cmap='gray', interpolation='nearest')
  im.set_clim([0, 1])
  plt.tight_layout(pad=0, w_pad=0, h_pad=0)
  # pylint: disable=invalid-name
  def update_img(pair):
    """Decode pairs of image strings, update a video."""
    im_i, im_j = pair
    nparr_i = np.fromstring(str(im_i), np.uint8)
    img_np_i = cv2.imdecode(nparr_i, 1)
    img_np_i = img_np_i[..., [2, 1, 0]]
    nparr_j = np.fromstring(str(im_j), np.uint8)
    img_np_j = cv2.imdecode(nparr_j, 1)
    img_np_j = img_np_j[..., [2, 1, 0]]

    # Optionally reshape the images to be same size.
    frame = np.concatenate([img_np_i, img_np_j], axis=1)
    im.set_data(frame)
    return im
  ani = animation.FuncAnimation(fig, update_img, combined, interval=15)
  writer = animation.writers['ffmpeg'](fps=15)
  dpi = 100
  tf.logging.info('Writing video to:\n %s \n' % vid_path)
  ani.save('%s.mp4' % vid_path, writer=writer, dpi=dpi) 

def save_fitted_as_movie(x_v, matv, results,
                         p1, p2, data_all, param_dict,
                         result_folder, prefix=None, use_snip=False, dpi=150):
    """
    Create movie to save single pixel fitting resutls.
    """
    total_n = data_all.shape[1]*p2[0]

    fig, ax = plt.subplots(nrows=1, ncols=1)
    ax.set_aspect('equal')
    ax.set_xlabel('Energy [keV]')
    ax.set_ylabel('Counts')
    max_v = np.max(data_all[p1[0]:p2[0], p1[1]:p2[1], :])
    ax.set_ylim([0, 1.1*max_v])

    l1,  = ax.plot(x_v,  x_v, label='exp', linestyle='-', marker='.')
    l2,  = ax.plot(x_v,  x_v, label='fit', color='red', linewidth=2)

    # fitted_sum = None
    plist = []
    for v in range(total_n):
        m = v / data_all.shape[1]
        n = v % data_all.shape[1]
        if m >= p1[0] and m <= p2[0] and n >= p1[1] and n <= p2[1]:
            plist.append((m, n))

    def update_img(p_val):
        m = p_val[0]
        n = p_val[1]
        data_y = data_all[m, n, :]

        fitted_y = np.sum(matv*results[m, n, :], axis=1)
        if use_snip is True:
            bg = snip_method_numba(data_y,
                                   param_dict['e_offset']['value'],
                                   param_dict['e_linear']['value'],
                                   param_dict['e_quadratic']['value'],
                                   width=param_dict['non_fitting_values']['background_width'])
            fitted_y += bg

        ax.set_title('Single pixel fitting for point ({}, {})'.format(m, n))
        # ax.set_ylim(low_limit_v, max_v*2)
        l1.set_ydata(data_y)
        l2.set_ydata(fitted_y)
        return l1, l2

    writer = animation.writers['ffmpeg'](fps=30)
    ani = animation.FuncAnimation(fig, update_img, plist)
    if prefix:
        output_file = prefix+'_pixel.mp4'
    else:
        output_file = 'fit_pixel.mp4'
    output_p = os.path.join(result_folder, output_file)
    ani.save(output_p, writer=writer, dpi=dpi) 

def create_movie(data, fname='demo.mp4', dpi=100, cmap='jet',
                 clim=None, fig_size=(6, 8), fps=20, data_power=1, angle=None, runid=None):
    """
    Transfer 3d array into a movie.

    Parameters
    ----------
    data : 3d array
        data shape is [num_sequences, num_row, num_col]
    fname : string, optional
        name to save movie
    dpi : int, optional
        resolution of the movie
    cmap : string, optional
        color format
    clim : list, tuple, optional
        [low, high] value to define plotting range
    fig_size : list, tuple, optional
        size (horizontal size, vertical size) of each plot
    fps : int, optional
        frame per second
    """
    fig, ax = plt.subplots()
    ax.set_aspect('equal')
    ax.get_xaxis().set_visible(False)
    ax.get_yaxis().set_visible(False)

    im = ax.imshow(np.zeros([data.shape[1], data.shape[2]]),
                   cmap=cmap, interpolation='nearest')

    fig.set_size_inches(fig_size)
    fig.tight_layout()

    def update_img(n):
        tmp = data[n, :, :]
        im.set_data(tmp**data_power)
        if clim is not None:
            im.set_clim(clim)
        else:
            im.set_clim([0, np.max(data[n, :, :])])
        figname = ''
        if runid is not None:
            figname = 'runid: {} '.format(runid[n])
        if angle is not None:
            figname += 'angle: {}'.format(angle[n])
        # if len(figname) != 0:
        #     im.ax.set_title(figname)
        return im

    # legend(loc=0)
    ani = animation.FuncAnimation(fig, update_img, data.shape[0], interval=30)
    writer = animation.writers['ffmpeg'](fps=fps)

    ani.save(fname, writer=writer, dpi=dpi) 

def evaluate_model(nb_traj,policy,save_dir= None,animate=False,**kwargs):
    transition_noise = kwargs['transition_noise'] if 'transition_noise' in kwargs.keys() else 0.02
    location_noise = kwargs['location_noise'] if 'location_noise' in kwargs.keys() else 0.1
    env = NavEnv(location_noise, transition_noise,done_thresh=0.03)
    if kwargs['zero_shot']:
        test_sketches = [[0, 1, 3,0], [3, 2, 3,2], [0, 2, 1,0], [1, 3, 0,2], [1, 3, 0,3], [1, 3, 2,1], [1, 3, 2,3], [2, 3, 2,1], [1, 2, 3,1],
                      [3, 1, 0,1], [3, 1, 2,0], [2, 3, 0,1],[1,3,1,3],[2,3,2,3],[1,2,1,2],[0,1,0,1],[0,2,0,2],[0,3,0,3]]
    else:
        test_sketches = [[0, 1, 3], [1, 2, 3], [3, 2, 1], [2, 3, 0], [1, 2, 0], [1, 0, 2], [1, 3, 2], [0, 3, 2],
                          [1, 2, 3], [3, 2, 0], [2, 1, 3], [1, 3, 0]]

    score = []
    if animate:
        fig = plt.figure()
        ims = []
    for i in range(nb_traj):
        task_score = []
        g_state,r_state = env.reset()
        colours = ['b', 'r', 'g', 'y', 'k']
        # randomly sample integets that constitute the sketch
        sketch = test_sketches[np.random.randint(0,len(test_sketches),1)[0]]
        curr_idx = 0
        curr_subtask = sketch[curr_idx]
        counter=0
        while True:
            action,stop = policy.forward_full([[r_state]], curr_subtask, dropout=1.)
            g_state,r_state,done = env.step(action,curr_subtask)
            if stop == 1 or counter >100:
                if done:
                    task_score.append(1)
                else:
                    task_score.append(0)
                if curr_idx < len(sketch) - 1:
                    curr_idx += 1
                    curr_subtask = sketch[curr_idx]
                    counter = 0
                else:
                    score.append(task_score)
                    break

            if animate:
                ims.append((plt.scatter(g_state[:, 0], g_state[:, 1], c=colours),))
            counter+=1
    if animate:
        print('writing video at:',save_dir+'im_mod.mp4')
        Writer = animation.writers['ffmpeg']
        writer = Writer(fps=15, metadata=dict(artist='Me'), bitrate=1800)
        im_ani = animation.ArtistAnimation(fig, ims, interval=50, repeat_delay=3000,
                                           blit=True)
        im_ani.save(save_dir+'im_mod.mp4', writer=writer)
    acc = 0
    for s in score:
        if np.sum(sum(s)) == len(s):
            acc+=1
    acc /=float(len(score))
    # also returns accuracy so you dont calculate it outside
    return score,acc 

def collect_data(num_traj,save_dir,animate=False,**kwargs):
    transition_noise = kwargs['transition_noise'] if 'transition_noise' in kwargs.keys() else 0.05
    location_noise = kwargs['location_noise'] if 'location_noise' in kwargs.keys() else 0.2
    env = NavEnv(location_noise,transition_noise)
    dataset = {'states':[],'actions':[],'gt_onsets':[],'tasks':[],'params':kwargs}
    train_sketches = [[0,1,3],[1,2,3],[3,2,1],[2,3,0],[1,2,0],[1,0,2],[1,3,2],[0,3,2],[1,2,3],[3,2,0],[2,1,3],[1,3,0]]

    if animate:
        fig = plt.figure()
        ims = []

    for i in range(num_traj):
        g_state,r_state = env.reset()
        colours = ['b', 'r', 'g', 'y', 'k']
        # randomly sample integets that constitute the sketch
        sketch = train_sketches[np.random.randint(0,len(train_sketches),1)[0]]
        curr_idx = 0
        curr_subtask = sketch[curr_idx]
        sketch_traj = []
        # begin trajectory
        traj_states = []
        traj_actions = []
        while True:
            sketch_traj.append(curr_subtask) # This gives us ground truth about the task executed at each timestep
            #all_pos = np.array([agent_pos,red_pos,green_pos,yel_pos,black_pos])
            #state = np.ravel(get_state(all_pos[0],all_pos[1:],type = 'rel'))
            traj_states.append(r_state)
            action = get_action(g_state,curr_subtask)
            g_state,r_state,done = env.step(action,curr_subtask)
            traj_actions.append(action)
            if animate:
                ims.append((plt.scatter(g_state[:, 0], g_state[:, 1], c=colours),))
            if done:
                if curr_idx<len(sketch)-1:
                    curr_idx+=1

                    curr_subtask = sketch[curr_idx]
                else:
                    dataset['states'].append(traj_states)
                    dataset['actions'].append(traj_actions)
                    dataset['gt_onsets'].append(sketch_traj)
                    dataset['tasks'].append(sketch)
                    break
    save_dir = save_dir if save_dir[-2:] == '.p' else save_dir+'.p'
    pickle.dump(dataset,open(save_dir,'wb'))

    if animate:
        print('WRITING')
        Writer = animation.writers['ffmpeg']
        writer = Writer(fps=15, metadata=dict(artist='Me'), bitrate=1800)
        im_ani = animation.ArtistAnimation(fig, ims, interval=50, repeat_delay=3000,
                                           blit=True)
        im_ani.save(save_dir+'im_mod.mp4', writer=writer)