from __future__ import print_function
import os
import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
def plot_reconstructions(data, recon_mean, loss, loss_type, epoch, args):
if args.input_type == 'multinomial':
# data is already between 0 and 1
num_classes = 256
# Find largest class logit
tmp = recon_mean.view(-1, num_classes, *args.input_size).max(dim=1)[1]
recon_mean = tmp.float() / (num_classes - 1.)
if epoch == 1:
if not os.path.exists(args.snap_dir + 'reconstruction/'):
os.makedirs(args.snap_dir + 'reconstruction/')
# VISUALIZATION: plot real images
plot_images(args, data.data.cpu().numpy()[0:9], args.snap_dir + 'reconstruction/', 'real', size_x=3, size_y=3)
# VISUALIZATION: plot reconstructions
if loss_type == 'bpd':
fname = str(epoch) + '_bpd_%5.3f' % loss
elif loss_type == 'elbo':
fname = str(epoch) + '_elbo_%6.4f' % loss
plot_images(args, recon_mean.data.cpu().numpy()[0:9], args.snap_dir + 'reconstruction/', fname, size_x=3, size_y=3)
def plot_images(args, x_sample, dir, file_name, size_x=3, size_y=3):
fig = plt.figure(figsize=(size_x, size_y))
# fig = plt.figure(1)
gs = gridspec.GridSpec(size_x, size_y)
gs.update(wspace=0.05, hspace=0.05)
for i, sample in enumerate(x_sample):
ax = plt.subplot(gs[i])
plt.axis('off')
ax.set_xticklabels([])
ax.set_yticklabels([])
ax.set_aspect('equal')
sample = sample.reshape((args.input_size[0], args.input_size[1], args.input_size[2]))
sample = sample.swapaxes(0, 2)
sample = sample.swapaxes(0, 1)
if (args.input_type == 'binary') or (args.input_type in ['multinomial'] and args.input_size[0] == 1):
sample = sample[:, :, 0]
plt.imshow(sample, cmap='gray', vmin=0, vmax=1)
else:
plt.imshow(sample)
plt.savefig(dir + file_name + '.png', bbox_inches='tight')
plt.close(fig)
