visualize.py 文件源码

def visualize(name):
    for counter, file in enumerate(sorted(glob.glob(os.path.join(args.train_dir, '{}_*.hy'.format(name))), key=os.path.getmtime)[-args.n:]):
        print (file)
        f = h5py.File(file, 'r')
        # I = np.zeros((args.h, args.num_frames * args.w, args.c))
        generated_frames = f[f.keys()[0]]
        _, _, h, w, c = generated_frames.shape
        h_low = (h - args.h) / 2
        h_high = (h + args.h) / 2
        w_low = (w - args.w) / 2
        w_high = (w + args.w) / 2
        # Take only first set of frames from batch
        II = []
        if args.c == 1:
            for j in range(args.num_frames):
                I = np.reshape(generated_frames[0, j, h_low:h_high, w_low:w_high, 0], (args.h, args.w))
                if (I < 1.0).all() and (I > -1.0).all():
                    print ('Image in [-1, 1]')
                    I = ((I + 1.0) / 2 * 255).astype(np.int32)

                II.append(I)

        else:
            for j in range(args.num_frames):
                I = np.reshape(generated_frames[0, j, h_low:h_high, w_low:w_high, 0:args.c], (args.h, args.w, args.c))
                II.append(I)

        # II = np.stack(II)
        output_img_path = './outputs/{}_{}_{}.png'.format(args.output_prefix, name, str(counter))
        print ('Writing image:', output_img_path)
        print (len(II), II[0].shape)
        imageio.mimwrite(output_img_path, II)