def make_reconstruction_movie(ims, titles=None, n_rows=0, n_cols=0, save_file=None, frame_rate=15, dpi=100): """Produce movie with original video and reconstructed videos. `ims` and `titles` are corresponding lists; this data is plotted using a linear index, i.e. if n_rows = 2 and n_cols = 3 the image stack in ims[2] will be in the first row, second column; the image stack in ims[4] will be in the second row, first column. If ims[i] is empty, that grid location will be skipped. Parameters ---------- ims : :obj:`list` of :obj:`np.ndarray` each list element is of shape (n_frames, n_channels, y_pix, x_pix) titles : :obj:`list` of :obj:`str`, optional title for each panel n_rows : :obj:`int` number of rows in video grid layout n_cols : :obj:`int` number of columns in video grid layout save_file : :obj:`str`, optional full save file (path and filename) frame_rate : :obj:`float`, optional frame rate of saved movie dpi : :obj:`int`, optional dpi of movie figure; controls resolution of titles """ for im in ims: if len(im) != 0: n_frames, n_channels, y_pix, x_pix = im.shape break scale_ = 5 fig_width = scale_ * n_cols * n_channels / 2 fig_height = y_pix / x_pix * scale_ * n_rows / 2 offset = 0.5 if n_rows == 1 else 0 fig = plt.figure(figsize=(fig_width, fig_height + offset), dpi=dpi) gs = GridSpec(n_rows, n_cols, figure=fig) axs = [] ax_count = 0 for i in range(n_rows): for j in range(n_cols): if ax_count < len(ims): axs.append(fig.add_subplot(gs[i, j])) ax_count += 1 else: break for ax_i, ax in enumerate(axs): ax.set_xticks([]) ax.set_yticks([]) if len(ims[ax_i]) == 0: ax.set_axis_off() fontsize = 12 titles = ['' for _ in range(n_cols * n_rows)] if titles is None else titles for ax_i, ax in enumerate(axs): if len(ims[ax_i]) != 0: ax.set_title(titles[ax_i], fontsize=fontsize) default_kwargs = {'animated': True, 'cmap': 'gray', 'vmin': 0, 'vmax': 1} # ims is a list of lists, each row is a list of artists to draw in the current frame; here we # are just animating one artist, the image, in each frame ims_ani = [] for i in range(n_frames): ims_curr = [] for ax_i, ax in enumerate(axs): if len(ims[ax_i]) != 0: ims_tmp = ims[ax_i][i, 0] if n_channels == 1 else concat( ims[ax_i][i]) im = ax.imshow(ims_tmp, **default_kwargs) [s.set_visible(False) for s in ax.spines.values()] ims_curr.append(im) ims_ani.append(ims_curr) plt.tight_layout(pad=0) ani = animation.ArtistAnimation(fig, ims_ani, blit=True, repeat_delay=1000) save_movie(save_file, ani, frame_rate=frame_rate)
def make_ae_reconstruction_movie(ims_orig, ims_recon_ae, ims_recon_lin=None, save_file=None, frame_rate=15): """Produce movie with original video, reconstructed video, and residual. Parameters ---------- ims_orig : :obj:`np.ndarray` shape (n_frames, n_channels, y_pix, x_pix) ims_recon_ae : :obj:`np.ndarray` shape (n_frames, n_channels, y_pix, x_pix) ims_recon_lin : :obj:`np.ndarray`, optional shape (n_frames, n_channels, y_pix, x_pix) save_file : :obj:`str`, optional full save file (path and filename) frame_rate : :obj:`float`, optional frame rate of saved movie """ n_frames, n_channels, y_pix, x_pix = ims_orig.shape n_cols = 1 if ims_recon_lin is None else 2 n_rows = 3 offset = 1 # 0 if ims_recon_lin is None else 1 scale_ = 5 fig_width = scale_ * n_cols * n_channels / 2 fig_height = y_pix / x_pix * scale_ * n_rows / 2 fig = plt.figure(figsize=(fig_width, fig_height + offset), dpi=100) gs = GridSpec(n_rows, n_cols, figure=fig) axs = [] if ims_recon_lin is not None: axs.append(fig.add_subplot(gs[0, 0])) # 0: original frames axs.append(fig.add_subplot(gs[1, 0])) # 1: ae reconstructed frames axs.append(fig.add_subplot(gs[1, 1])) # 2: ae residuals axs.append(fig.add_subplot(gs[2, 0])) # 3: linear reconstructed frames axs.append(fig.add_subplot(gs[2, 1])) # 4: linear residuals else: axs.append(fig.add_subplot(gs[0, 0])) # 0: original frames axs.append(fig.add_subplot(gs[1, 0])) # 1: ae reconstructed frames axs.append(fig.add_subplot(gs[2, 0])) # 2: ae residuals for ax in fig.axes: ax.set_xticks([]) ax.set_yticks([]) fontsize = 12 axs[0].set_title('Original', fontsize=fontsize) axs[1].set_title('Conv AE reconstructed', fontsize=fontsize) axs[2].set_title('Conv AE residual', fontsize=fontsize) if ims_recon_lin is not None: axs[3].set_title('Linear AE reconstructed', fontsize=fontsize) axs[4].set_title('Linear AE residual', fontsize=fontsize) ims_res_ae = ims_orig - ims_recon_ae if ims_recon_lin is not None: ims_res_lin = ims_orig - ims_recon_lin else: ims_res_lin = None default_kwargs = {'animated': True, 'cmap': 'gray', 'vmin': 0, 'vmax': 1} # ims is a list of lists, each row is a list of artists to draw in the current frame; here we # are just animating one artist, the image, in each frame ims = [] for i in range(ims_orig.shape[0]): ims_curr = [] # original video ims_tmp = ims_orig[i, 0] if n_channels == 1 else concat(ims_orig[i]) im = axs[0].imshow(ims_tmp, **default_kwargs) [s.set_visible(False) for s in axs[0].spines.values()] ims_curr.append(im) # ae reconstructed video ims_tmp = ims_recon_ae[i, 0] if n_channels == 1 else concat( ims_recon_ae[i]) im = axs[1].imshow(ims_tmp, **default_kwargs) [s.set_visible(False) for s in axs[1].spines.values()] ims_curr.append(im) # ae residual video ims_tmp = ims_res_ae[i, 0] if n_channels == 1 else concat(ims_res_ae[i]) im = axs[2].imshow(0.5 + ims_tmp, **default_kwargs) [s.set_visible(False) for s in axs[2].spines.values()] ims_curr.append(im) if ims_recon_lin is not None: # linear reconstructed video ims_tmp = ims_recon_lin[i, 0] if n_channels == 1 else concat( ims_recon_lin[i]) im = axs[3].imshow(ims_tmp, **default_kwargs) [s.set_visible(False) for s in axs[3].spines.values()] ims_curr.append(im) # linear residual video ims_tmp = ims_res_lin[i, 0] if n_channels == 1 else concat( ims_res_lin[i]) im = axs[4].imshow(0.5 + ims_tmp, **default_kwargs) [s.set_visible(False) for s in axs[4].spines.values()] ims_curr.append(im) ims.append(ims_curr) plt.tight_layout(pad=0) ani = animation.ArtistAnimation(fig, ims, blit=True, repeat_delay=1000) writer = FFMpegWriter(fps=frame_rate, bitrate=-1) if save_file is not None: make_dir_if_not_exists(save_file) if save_file[-3:] != 'mp4': save_file += '.mp4' print('saving video to %s...' % save_file, end='') ani.save(save_file, writer=writer) # if save_file[-3:] != 'gif': # save_file += '.gif' # ani.save(save_file, writer='imagemagick', fps=15) print('done')
def make_neural_reconstruction_movie(ims_orig, ims_recon_ae, ims_recon_neural, latents_ae, latents_neural, ae_model_class='AE', colored_predictions=False, scale=0.5, xtick_locs=None, frame_rate_beh=None, save_file=None, frame_rate=15): """Produce movie with original video, ae reconstructed video, and neural reconstructed video. Latent traces are additionally plotted, as well as the residual between the ae reconstruction and the neural reconstruction. Parameters ---------- ims_orig : :obj:`np.ndarray` original images; shape (n_frames, n_channels, y_pix, x_pix) ims_recon_ae : :obj:`np.ndarray` images reconstructed by AE; shape (n_frames, n_channels, y_pix, x_pix) ims_recon_neural : :obj:`np.ndarray` images reconstructed by neural activity; shape (n_frames, n_channels, y_pix, x_pix) latents_ae : :obj:`np.ndarray` original AE latents; shape (n_frames, n_latents) latents_neural : :obj:`np.ndarray` latents reconstruted by neural activity; shape (n_frames, n_latents) ae_model_class : :obj:`str`, optional 'AE', 'VAE', etc. for plot titles colored_predictions : :obj:`bool`, optional False to plot reconstructions in black, True to plot in different colors scale : :obj:`int`, optional scale magnitude of traces xtick_locs : :obj:`array-like`, optional tick locations in units of bins frame_rate_beh : :obj:`float`, optional frame rate of behavorial video; to properly relabel xticks save_file : :obj:`str`, optional full save file (path and filename) frame_rate : :obj:`float`, optional frame rate of saved movie """ means = np.nanmean(latents_ae, axis=0) std = np.nanstd(latents_ae) / scale latents_ae_sc = (latents_ae - means) / std latents_dec_sc = (latents_neural - means) / std n_channels, y_pix, x_pix = ims_orig.shape[1:] n_time, n_ae_latents = latents_ae.shape n_cols = 3 n_rows = 2 offset = 2 # 0 if ims_recon_lin is None else 1 scale_ = 5 fig_width = scale_ * n_cols * n_channels / 2 fig_height = y_pix / x_pix * scale_ * n_rows / 2 fig = plt.figure(figsize=(fig_width, fig_height + offset)) gs = GridSpec(n_rows, n_cols, figure=fig) axs = [] axs.append(fig.add_subplot(gs[0, 0])) # 0: original frames axs.append(fig.add_subplot(gs[0, 1])) # 1: ae reconstructed frames axs.append(fig.add_subplot(gs[0, 2])) # 2: neural reconstructed frames axs.append(fig.add_subplot(gs[1, 0])) # 3: residual axs.append(fig.add_subplot(gs[1, 1:3])) # 4: ae and predicted ae latents for i, ax in enumerate(fig.axes): ax.set_yticks([]) if i > 2: ax.get_xaxis().set_tick_params(labelsize=12, direction='in') axs[0].set_xticks([]) axs[1].set_xticks([]) axs[2].set_xticks([]) axs[3].set_xticks([]) # check that the axes are correct fontsize = 12 idx = 0 axs[idx].set_title('Original', fontsize=fontsize) idx += 1 axs[idx].set_title('%s reconstructed' % ae_model_class, fontsize=fontsize) idx += 1 axs[idx].set_title('Neural reconstructed', fontsize=fontsize) idx += 1 axs[idx].set_title('Reconstructions residual', fontsize=fontsize) idx += 1 axs[idx].set_title('%s latent predictions' % ae_model_class, fontsize=fontsize) if xtick_locs is not None and frame_rate_beh is not None: axs[idx].set_xticks(xtick_locs) axs[idx].set_xticklabels( (np.asarray(xtick_locs) / frame_rate_beh).astype('int')) axs[idx].set_xlabel('Time (s)', fontsize=fontsize) else: axs[idx].set_xlabel('Time (bins)', fontsize=fontsize) time = np.arange(n_time) ims_res = ims_recon_ae - ims_recon_neural im_kwargs = {'animated': True, 'cmap': 'gray', 'vmin': 0, 'vmax': 1} tr_kwargs = {'animated': True, 'linewidth': 2} latents_ae_color = [0.2, 0.2, 0.2] label_ae_base = '%s latents' % ae_model_class label_dec_base = 'Predicted %s latents' % ae_model_class # ims is a list of lists, each row is a list of artists to draw in the # current frame; here we are just animating one artist, the image, in # each frame ims = [] for i in range(n_time): ims_curr = [] idx = 0 if i % 100 == 0: print('processing frame %03i/%03i' % (i, n_time)) ################### # behavioral videos ################### # original video ims_tmp = ims_orig[i, 0] if n_channels == 1 else concat(ims_orig[i]) im = axs[idx].imshow(ims_tmp, **im_kwargs) ims_curr.append(im) idx += 1 # ae reconstruction ims_tmp = ims_recon_ae[i, 0] if n_channels == 1 else concat( ims_recon_ae[i]) im = axs[idx].imshow(ims_tmp, **im_kwargs) ims_curr.append(im) idx += 1 # neural reconstruction ims_tmp = ims_recon_neural[i, 0] if n_channels == 1 else concat( ims_recon_neural[i]) im = axs[idx].imshow(ims_tmp, **im_kwargs) ims_curr.append(im) idx += 1 # residual ims_tmp = ims_res[i, 0] if n_channels == 1 else concat(ims_res[i]) im = axs[idx].imshow(0.5 + ims_tmp, **im_kwargs) ims_curr.append(im) idx += 1 ######## # traces ######## # latents over time axs[idx].set_prop_cycle(None) # reset colors for latent in range(n_ae_latents): if colored_predictions: latents_dec_color = axs[idx]._get_lines.get_next_color() else: latents_dec_color = [0, 0, 0] # just put labels on last lvs if latent == n_ae_latents - 1 and i == 0: label_ae = label_ae_base label_dec = label_dec_base else: label_ae = None label_dec = None im = axs[idx].plot(time[0:i + 1], latent + latents_ae_sc[0:i + 1, latent], color=latents_ae_color, alpha=0.7, label=label_ae, **tr_kwargs)[0] axs[idx].spines['top'].set_visible(False) axs[idx].spines['right'].set_visible(False) axs[idx].spines['left'].set_visible(False) ims_curr.append(im) im = axs[idx].plot(time[0:i + 1], latent + latents_dec_sc[0:i + 1, latent], color=latents_dec_color, label=label_dec, **tr_kwargs)[0] axs[idx].spines['top'].set_visible(False) axs[idx].spines['right'].set_visible(False) axs[idx].spines['left'].set_visible(False) if colored_predictions: # original latents - gray orig_line = mlines.Line2D([], [], color=[0.2, 0.2, 0.2], linewidth=3, alpha=0.7) # predicted latents - cycle through some colors colors = plt.rcParams['axes.prop_cycle'].by_key()['color'] dls = [] for c in range(5): dls.append( mlines.Line2D([], [], linewidth=3, linestyle='--', dashes=(0, 3 * c, 20, 1), color='%s' % colors[c])) plt.legend([orig_line, tuple(dls)], [label_ae_base, label_dec_base], loc='lower right', fontsize=fontsize, frameon=True, framealpha=0.7, edgecolor=[1, 1, 1]) else: plt.legend(loc='lower right', fontsize=fontsize, frameon=True, framealpha=0.7, edgecolor=[1, 1, 1]) ims_curr.append(im) ims.append(ims_curr) plt.tight_layout(pad=0) ani = animation.ArtistAnimation(fig, ims, blit=True, repeat_delay=1000) save_movie(save_file, ani, frame_rate=frame_rate)
def make_neural_reconstruction_movie(ims_orig, ims_recon_ae, ims_recon_neural, latents_ae, latents_neural, save_file=None, frame_rate=15): """Produce movie with original video, ae reconstructed video, and neural reconstructed video. Latent traces are additionally plotted, as well as the residual between the ae reconstruction and the neural reconstruction. Parameters ---------- ims_orig : :obj:`np.ndarray` shape (n_frames, n_channels, y_pix, x_pix) ims_recon_ae : :obj:`np.ndarray` shape (n_frames, n_channels, y_pix, x_pix) ims_recon_neural : :obj:`np.ndarray`, optional shape (n_frames, n_channels, y_pix, x_pix) latents_ae : :obj:`np.ndarray`, optional shape (n_frames, n_latents) save_file : :obj:`str`, optional full save file (path and filename) frame_rate : :obj:`float`, optional frame rate of saved movie """ means = np.mean(latents_ae, axis=0) std = np.std(latents_ae) * 2 latents_ae_sc = (latents_ae - means) / std latents_dec_sc = (latents_neural - means) / std n_channels, y_pix, x_pix = ims_orig.shape[1:] n_time, n_ae_latents = latents_ae.shape n_cols = 3 n_rows = 2 offset = 2 # 0 if ims_recon_lin is None else 1 scale_ = 5 fig_width = scale_ * n_cols * n_channels / 2 fig_height = y_pix / x_pix * scale_ * n_rows / 2 fig = plt.figure(figsize=(fig_width, fig_height + offset)) gs = GridSpec(n_rows, n_cols, figure=fig) axs = [] axs.append(fig.add_subplot(gs[0, 0])) # 0: original frames axs.append(fig.add_subplot(gs[0, 1])) # 1: ae reconstructed frames axs.append(fig.add_subplot(gs[0, 2])) # 2: neural reconstructed frames axs.append(fig.add_subplot(gs[1, 0])) # 3: residual axs.append(fig.add_subplot(gs[1, 1:3])) # 4: ae and predicted ae latents for i, ax in enumerate(fig.axes): ax.set_yticks([]) if i > 2: ax.get_xaxis().set_tick_params(labelsize=12, direction='in') axs[0].set_xticks([]) axs[1].set_xticks([]) axs[2].set_xticks([]) axs[3].set_xticks([]) # check that the axes are correct fontsize = 12 idx = 0 axs[idx].set_title('Original', fontsize=fontsize) idx += 1 axs[idx].set_title('AE reconstructed', fontsize=fontsize) idx += 1 axs[idx].set_title('Neural reconstructed', fontsize=fontsize) idx += 1 axs[idx].set_title('Reconstructions residual', fontsize=fontsize) idx += 1 axs[idx].set_title('AE latent predictions', fontsize=fontsize) axs[idx].set_xlabel('Time (bins)', fontsize=fontsize) time = np.arange(n_time) ims_res = ims_recon_ae - ims_recon_neural im_kwargs = {'animated': True, 'cmap': 'gray', 'vmin': 0, 'vmax': 1} tr_kwargs = {'animated': True, 'linewidth': 2} latents_ae_color = [0.2, 0.2, 0.2] latents_dec_color = [0, 0, 0] # ims is a list of lists, each row is a list of artists to draw in the # current frame; here we are just animating one artist, the image, in # each frame ims = [] for i in range(n_time): ims_curr = [] idx = 0 if i % 100 == 0: print('processing frame %03i/%03i' % (i, n_time)) ################### # behavioral videos ################### # original video ims_tmp = ims_orig[i, 0] if n_channels == 1 else concat(ims_orig[i]) im = axs[idx].imshow(ims_tmp, **im_kwargs) ims_curr.append(im) idx += 1 # ae reconstruction ims_tmp = ims_recon_ae[i, 0] if n_channels == 1 else concat( ims_recon_ae[i]) im = axs[idx].imshow(ims_tmp, **im_kwargs) ims_curr.append(im) idx += 1 # neural reconstruction ims_tmp = ims_recon_neural[i, 0] if n_channels == 1 else concat( ims_recon_neural[i]) im = axs[idx].imshow(ims_tmp, **im_kwargs) ims_curr.append(im) idx += 1 # residual ims_tmp = ims_res[i, 0] if n_channels == 1 else concat(ims_res[i]) im = axs[idx].imshow(0.5 + ims_tmp, **im_kwargs) ims_curr.append(im) idx += 1 ######## # traces ######## # latents over time for latent in range(n_ae_latents): # just put labels on last lvs if latent == n_ae_latents - 1 and i == 0: label_ae = 'AE latents' label_dec = 'Predicted AE latents' else: label_ae = None label_dec = None im = axs[idx].plot(time[0:i + 1], latent + latents_ae_sc[0:i + 1, latent], color=latents_ae_color, alpha=0.7, label=label_ae, **tr_kwargs)[0] axs[idx].spines['top'].set_visible(False) axs[idx].spines['right'].set_visible(False) axs[idx].spines['left'].set_visible(False) ims_curr.append(im) im = axs[idx].plot(time[0:i + 1], latent + latents_dec_sc[0:i + 1, latent], color=latents_dec_color, label=label_dec, **tr_kwargs)[0] axs[idx].spines['top'].set_visible(False) axs[idx].spines['right'].set_visible(False) axs[idx].spines['left'].set_visible(False) plt.legend(loc='lower right', fontsize=fontsize, frameon=True, framealpha=0.7, edgecolor=[1, 1, 1]) ims_curr.append(im) ims.append(ims_curr) plt.tight_layout(pad=0) ani = animation.ArtistAnimation(fig, ims, blit=True, repeat_delay=1000) writer = FFMpegWriter(fps=frame_rate, bitrate=-1) if save_file is not None: make_dir_if_not_exists(save_file) if save_file[-3:] != 'mp4': save_file += '.mp4' print('saving video to %s...' % save_file, end='') ani.save(save_file, writer=writer) print('done')