Update app.py

app.py

@@ -610,7 +610,6 @@ def name2seq(file_name):
 
     images = np.asarray(images, dtype='f') / 256.0
     images = images.transpose((0, 3, 1, 2))
-    print(images.shape)
     images = torch.Tensor(images).unsqueeze(dim=0)
     return images
     
@@ -649,6 +648,58 @@ def display_image(file_name):
     imageio.imwrite('image.png', image)
     
 
+def concat(file_name):
+    images = []
+
+    for frame in range(8):
+        frame_name = '%d' % (frame)
+        image_filename = file_name + frame_name + '.png'
+        image = imageio.imread(image_filename)
+        images.append(image)
+
+    gif_filename = 'demo.gif'
+    return imageio.mimsave(gif_filename, images)
+    
+
+def MyPlot(frame_id, src_orig, tar_orig, src_recon, tar_recon, src_Zt, tar_Zt, src_Zf_tar_Zt, tar_Zf_src_Zt):
+    
+    fig, axs = plt.subplots(2, 4, sharex=True, sharey=True, figsize=(10, 5))
+    
+    axs[0, 0].imshow(src_orig)
+    axs[0, 0].set_title("\n\n\nOriginal\nInput")
+    axs[0, 0].axis('off')
+
+    axs[1, 0].imshow(tar_orig)
+    axs[1, 0].axis('off')
+
+    axs[0, 1].imshow(src_recon)
+    axs[0, 1].set_title("\n\n\nReconstructed\nOutput")
+    axs[0, 1].axis('off')
+
+    axs[1, 1].imshow(tar_recon)
+    axs[1, 1].axis('off')
+
+    axs[0, 2].imshow(src_Zt)
+    axs[0, 2].set_title("\n\n\nOutput\nw/ Zt")
+    axs[0, 2].axis('off')
+
+    axs[1, 2].imshow(tar_Zt)
+    axs[1, 2].axis('off')
+    
+    axs[0, 3].imshow(tar_Zf_src_Zt)
+    axs[0, 3].set_title("\n\n\nExchange\nZt and Zf")
+    axs[0, 3].axis('off')
+
+    axs[1, 3].imshow(src_Zf_tar_Zt)
+    axs[1, 3].axis('off')
+    
+    plt.subplots_adjust(hspace=0.06, wspace=0.05)
+    
+    save_name = 'MyPlot_{}.png'.format(frame_id)
+    
+    plt.savefig(save_name, dpi=200, format='png', bbox_inches='tight', pad_inches=0.0)
+    
+
 def run(domain_source, action_source, hair_source, top_source, bottom_source, domain_target, action_target, hair_target, top_target, bottom_target):
 
     # == Source Avatar ==
@@ -676,7 +727,6 @@ def run(domain_source, action_source, hair_source, top_source, bottom_source, do
     file_name_source = './Sprite/frames/domain_1/' + action_source + '/'
     file_name_source = file_name_source + 'front' + '_' + str(body_source) + str(bottom_source) + str(top_source) + str(hair_source) + '_'
     
-    gif = display_gif_pad(file_name_source, 'avatar_source.gif')
     
     # == Target Avatar ==
     # body
@@ -703,7 +753,10 @@ def run(domain_source, action_source, hair_source, top_source, bottom_source, do
     file_name_target = './Sprite/frames/domain_2/' + action_target + '/'
     file_name_target = file_name_target + 'front' + '_' + str(body_target) + str(bottom_target) + str(top_target) + str(hair_target) + '_'
     
-    gif_target = display_gif_pad(file_name_target, 'avatar_target.gif')
+    
+    # == Load Input ==
+    images_source = name2seq(file_name_source)
+    images_target = name2seq(file_name_target)
     
     
     # == Load Model ==
@@ -711,6 +764,59 @@ def run(domain_source, action_source, hair_source, top_source, bottom_source, do
     model.load_state_dict(torch.load('TransferVAE.pth.tar', map_location=torch.device('cpu'))['state_dict'])
     model.eval()
     
+    
+    # == Forward ==
+    with torch.no_grad():
+        f_mean, f_logvar, f_post, z_post_mean, z_post_logvar, z_post, z_prior_mean, z_prior_logvar, z_prior, recon_x, pred_domain_all, pred_video_class = model(x, [0]*3)
+     
+    src_orig_sample = x[0, :, :, :, :]
+    src_recon_sample = recon_x[0, :, :, :, :]
+    src_f_post = f_post[0, :].unsqueeze(0)
+    src_z_post = z_post[0, :, :].unsqueeze(0)
+
+    tar_orig_sample = x[1, :, :, :, :]
+    tar_recon_sample = recon_x[1, :, :, :, :]
+    tar_f_post = f_post[1, :].unsqueeze(0)
+    tar_z_post = z_post[1, :, :].unsqueeze(0)   
+    
+    
+    # == Visualize ==
+    for frame in range(8):
+    
+        # original frame
+        src_orig = src_orig_sample[frame, :, :, :].detach().numpy().transpose((1, 2, 0))
+        tar_orig = tar_orig_sample[frame, :, :, :].detach().numpy().transpose((1, 2, 0))
+    
+        # reconstructed frame
+        src_recon = src_recon_sample[frame, :, :, :].detach().numpy().transpose((1, 2, 0))
+        tar_recon = tar_recon_sample[frame, :, :, :].detach().numpy().transpose((1, 2, 0))
+    
+        # Zt
+        f_expand_src = 0 * src_f_post.unsqueeze(1).expand(-1, 8, opt.f_dim)
+        zf_src = torch.cat((src_z_post, f_expand_src), dim=2)
+        recon_x_src = model.decoder_frame(zf_src)
+        src_Zt = recon_x_src.squeeze()[frame, :, :, :].detach().numpy().transpose((1, 2, 0))
+    
+        f_expand_tar = 0 * tar_f_post.unsqueeze(1).expand(-1, 8, opt.f_dim)
+        zf_tar = torch.cat((tar_z_post, f_expand_tar), dim=2) # batch,frames,(z_dim+f_dim)
+        recon_x_tar = model.decoder_frame(zf_tar)
+        tar_Zt = recon_x_tar.squeeze()[frame, :, :, :].detach().numpy().transpose((1, 2, 0))
+    
+        # Zf_Zt
+        f_expand_src = src_f_post.unsqueeze(1).expand(-1, 8, opt.f_dim)
+        zf_srcZf_tarZt = torch.cat((tar_z_post, f_expand_src), dim=2) # batch,frames,(z_dim+f_dim)
+        recon_x_srcZf_tarZt = model.decoder_frame(zf_srcZf_tarZt)
+        src_Zf_tar_Zt = recon_x_srcZf_tarZt.squeeze()[frame, :, :, :].detach().numpy().transpose((1, 2, 0))
+    
+        f_expand_tar = tar_f_post.unsqueeze(1).expand(-1, 8, opt.f_dim)
+        zf_tarZf_srcZt = torch.cat((src_z_post, f_expand_tar), dim=2) # batch,frames,(z_dim+f_dim)
+        recon_x_tarZf_srcZt = model.decoder_frame(zf_tarZf_srcZt)
+        tar_Zf_src_Zt = recon_x_tarZf_srcZt.squeeze()[frame, :, :, :].detach().numpy().transpose((1, 2, 0))
+
+        MyPlot(frame, src_orig, tar_orig, src_recon, tar_recon, src_Zt, tar_Zt, src_Zf_tar_Zt, tar_Zf_src_Zt)
+        
+    a = concat('MyPlot_')
+    
     return 'demo.gif'