visual_evaluation.py 文件源码

def plot_reconstruction( args, samples, c, dir , size_x=3, size_y=3):
    samples = samples.data.cpu().numpy()[:size_x * size_y]

    fig = plt.figure(figsize=(size_x, size_y))
    gs = gridspec.GridSpec(size_x, size_y)
    gs.update(wspace=0.05, hspace=0.05)

    for i, sample in enumerate(samples):
        ax = plt.subplot(gs[i])
        plt.axis('off')
        ax.set_xticklabels([])
        ax.set_yticklabels([])
        ax.set_aspect('equal')
        plt.imshow(sample.reshape(args.input_size[1], args.input_size[2]), cmap='Greys_r')

    if not os.path.exists(dir + 'reconstruction/'):
        os.makedirs(dir + 'reconstruction/')

    plt.savefig(dir + 'reconstruction/{}.png'.format(str(c).zfill(3)), bbox_inches='tight')
    plt.close(fig)

#=======================================================================================================================