import gradio as gr import cv2 import imageio import math from math import ceil import matplotlib.pyplot as plt import matplotlib.animation as animation import numpy as np from PIL import Image import torch import torch.nn as nn import torch.nn.functional as F from torch.autograd import Function class RelationModuleMultiScale(torch.nn.Module): def __init__(self, img_feature_dim, num_bottleneck, num_frames): super(RelationModuleMultiScale, self).__init__() self.subsample_num = 3 self.img_feature_dim = img_feature_dim self.scales = [i for i in range(num_frames, 1, -1)] self.relations_scales = [] self.subsample_scales = [] for scale in self.scales: relations_scale = self.return_relationset(num_frames, scale) self.relations_scales.append(relations_scale) self.subsample_scales.append(min(self.subsample_num, len(relations_scale))) self.num_frames = num_frames self.fc_fusion_scales = nn.ModuleList() for i in range(len(self.scales)): scale = self.scales[i] fc_fusion = nn.Sequential(nn.ReLU(), nn.Linear(scale * self.img_feature_dim, num_bottleneck), nn.ReLU()) self.fc_fusion_scales += [fc_fusion] def forward(self, input): act_scale_1 = input[:, self.relations_scales[0][0] , :] act_scale_1 = act_scale_1.view(act_scale_1.size(0), self.scales[0] * self.img_feature_dim) act_scale_1 = self.fc_fusion_scales[0](act_scale_1) act_scale_1 = act_scale_1.unsqueeze(1) act_all = act_scale_1.clone() for scaleID in range(1, len(self.scales)): act_relation_all = torch.zeros_like(act_scale_1) num_total_relations = len(self.relations_scales[scaleID]) num_select_relations = self.subsample_scales[scaleID] idx_relations_evensample = [int(ceil(i * num_total_relations / num_select_relations)) for i in range(num_select_relations)] for idx in idx_relations_evensample: act_relation = input[:, self.relations_scales[scaleID][idx], :] act_relation = act_relation.view(act_relation.size(0), self.scales[scaleID] * self.img_feature_dim) act_relation = self.fc_fusion_scales[scaleID](act_relation) act_relation = act_relation.unsqueeze(1) act_relation_all += act_relation act_all = torch.cat((act_all, act_relation_all), 1) return act_all def return_relationset(self, num_frames, num_frames_relation): import itertools return list(itertools.combinations([i for i in range(num_frames)], num_frames_relation)) class GradReverse(Function): @staticmethod def forward(ctx, x, beta): ctx.beta = beta return x.view_as(x) @staticmethod def backward(ctx, grad_output): grad_input = grad_output.neg() * ctx.beta return grad_input, None class TransferVAE_Video(nn.Module): def __init__(self): super(TransferVAE_Video, self).__init__() self.f_dim = 512 self.z_dim = 512 self.fc_dim = 1024 self.channels = 3 self.frames = 8 self.batch_size = 128 self.dropout_rate = 0.5 self.num_class = 15 self.prior_sample = 'random' import dcgan_64 self.encoder = dcgan_64.encoder(self.fc_dim, self.channels) self.decoder = dcgan_64.decoder_woSkip(self.z_dim + self.f_dim, self.channels) self.fc_output_dim = self.fc_dim self.relu = nn.LeakyReLU(0.1) self.dropout_f = nn.Dropout(p=self.dropout_rate) self.dropout_v = nn.Dropout(p=self.dropout_rate) self.hidden_dim = 512 self.f_rnn_layers = 1 self.z_prior_lstm_ly1 = nn.LSTMCell(self.z_dim, self.hidden_dim) self.z_prior_lstm_ly2 = nn.LSTMCell(self.hidden_dim, self.hidden_dim) self.z_prior_mean = nn.Linear(self.hidden_dim, self.z_dim) self.z_prior_logvar = nn.Linear(self.hidden_dim, self.z_dim) self.z_lstm = nn.LSTM(self.fc_output_dim, self.hidden_dim, self.f_rnn_layers, bidirectional=True, batch_first=True) self.f_mean = nn.Linear(self.hidden_dim * 2, self.f_dim) self.f_logvar = nn.Linear(self.hidden_dim * 2, self.f_dim) self.z_rnn = nn.RNN(self.hidden_dim * 2, self.hidden_dim, batch_first=True) self.z_mean = nn.Linear(self.hidden_dim, self.z_dim) self.z_logvar = nn.Linear(self.hidden_dim, self.z_dim) self.fc_feature_domain_frame = nn.Linear(self.z_dim, self.z_dim) self.fc_classifier_domain_frame = nn.Linear(self.z_dim, 2) self.num_bottleneck = 256 self.TRN = RelationModuleMultiScale(self.z_dim, self.num_bottleneck, self.frames) self.bn_trn_S = nn.BatchNorm1d(self.num_bottleneck) self.bn_trn_T = nn.BatchNorm1d(self.num_bottleneck) self.feat_aggregated_dim = self.num_bottleneck self.fc_feature_domain_video = nn.Linear(self.feat_aggregated_dim, self.feat_aggregated_dim) self.fc_classifier_domain_video = nn.Linear(self.feat_aggregated_dim, 2) self.relation_domain_classifier_all = nn.ModuleList() for i in range(self.frames-1): relation_domain_classifier = nn.Sequential( nn.Linear(self.feat_aggregated_dim, self.feat_aggregated_dim), nn.ReLU(), nn.Linear(self.feat_aggregated_dim, 2) ) self.relation_domain_classifier_all += [relation_domain_classifier] self.pred_classifier_video = nn.Linear(self.feat_aggregated_dim, self.num_class) self.fc_feature_domain_latent = nn.Linear(self.f_dim, self.f_dim) self.fc_classifier_doamin_latent = nn.Linear(self.f_dim, 2) def domain_classifier_frame(self, feat, beta): feat_fc_domain_frame = GradReverse.apply(feat, beta) feat_fc_domain_frame = self.fc_feature_domain_frame(feat_fc_domain_frame) feat_fc_domain_frame = self.relu(feat_fc_domain_frame) pred_fc_domain_frame = self.fc_classifier_domain_frame(feat_fc_domain_frame) return pred_fc_domain_frame def domain_classifier_video(self, feat_video, beta): feat_fc_domain_video = GradReverse.apply(feat_video, beta) feat_fc_domain_video = self.fc_feature_domain_video(feat_fc_domain_video) feat_fc_domain_video = self.relu(feat_fc_domain_video) pred_fc_domain_video = self.fc_classifier_domain_video(feat_fc_domain_video) return pred_fc_domain_video def domain_classifier_latent(self, f): feat_fc_domain_latent = self.fc_feature_domain_latent(f) feat_fc_domain_latent = self.relu(feat_fc_domain_latent) pred_fc_domain_latent = self.fc_classifier_doamin_latent(feat_fc_domain_latent) return pred_fc_domain_latent def domain_classifier_relation(self, feat_relation, beta): pred_fc_domain_relation_video = None for i in range(len(self.relation_domain_classifier_all)): feat_relation_single = feat_relation[:,i,:].squeeze(1) feat_fc_domain_relation_single = GradReverse.apply(feat_relation_single, beta) pred_fc_domain_relation_single = self.relation_domain_classifier_all[i](feat_fc_domain_relation_single) if pred_fc_domain_relation_video is None: pred_fc_domain_relation_video = pred_fc_domain_relation_single.view(-1,1,2) else: pred_fc_domain_relation_video = torch.cat((pred_fc_domain_relation_video, pred_fc_domain_relation_single.view(-1,1,2)), 1) pred_fc_domain_relation_video = pred_fc_domain_relation_video.view(-1,2) return pred_fc_domain_relation_video def get_trans_attn(self, pred_domain): softmax = nn.Softmax(dim=1) logsoftmax = nn.LogSoftmax(dim=1) entropy = torch.sum(-softmax(pred_domain) * logsoftmax(pred_domain), 1) weights = 1 - entropy return weights def get_general_attn(self, feat): num_segments = feat.size()[1] feat = feat.view(-1, feat.size()[-1]) # reshape features: 128x4x256 --> (128x4)x256 weights = self.attn_layer(feat) # e.g. (128x4)x1 weights = weights.view(-1, num_segments, weights.size()[-1]) # reshape attention weights: (128x4)x1 --> 128x4x1 weights = F.softmax(weights, dim=1) # softmax over segments ==> 128x4x1 return weights def get_attn_feat_relation(self, feat_fc, pred_domain, num_segments): weights_attn = self.get_trans_attn(pred_domain) weights_attn = weights_attn.view(-1, num_segments-1, 1).repeat(1,1,feat_fc.size()[-1]) # reshape & repeat weights (e.g. 16 x 4 x 256) feat_fc_attn = (weights_attn+1) * feat_fc return feat_fc_attn, weights_attn[:,:,0] def encode_and_sample_post(self, x): if isinstance(x, list): conv_x = self.encoder_frame(x[0]) else: conv_x = self.encoder_frame(x) # pass the bidirectional lstm lstm_out, _ = self.z_lstm(conv_x) # get f: backward = lstm_out[:, 0, self.hidden_dim:2 * self.hidden_dim] frontal = lstm_out[:, self.frames - 1, 0:self.hidden_dim] lstm_out_f = torch.cat((frontal, backward), dim=1) f_mean = self.f_mean(lstm_out_f) f_logvar = self.f_logvar(lstm_out_f) f_post = self.reparameterize(f_mean, f_logvar, random_sampling=False) # pass to one direction rnn features, _ = self.z_rnn(lstm_out) z_mean = self.z_mean(features) z_logvar = self.z_logvar(features) z_post = self.reparameterize(z_mean, z_logvar, random_sampling=False) if isinstance(x, list): f_mean_list = [f_mean] f_post_list = [f_post] for t in range(1,3,1): conv_x = self.encoder_frame(x[t]) lstm_out, _ = self.z_lstm(conv_x) # get f: backward = lstm_out[:, 0, self.hidden_dim:2 * self.hidden_dim] frontal = lstm_out[:, self.frames - 1, 0:self.hidden_dim] lstm_out_f = torch.cat((frontal, backward), dim=1) f_mean = self.f_mean(lstm_out_f) f_logvar = self.f_logvar(lstm_out_f) f_post = self.reparameterize(f_mean, f_logvar, random_sampling=False) f_mean_list.append(f_mean) f_post_list.append(f_post) f_mean = f_mean_list f_post = f_post_list # f_mean and f_post are list if triple else not return f_mean, f_logvar, f_post, z_mean, z_logvar, z_post def decoder_frame(self,zf): recon_x = self.decoder(zf) return recon_x def encoder_frame(self, x): x_shape = x.shape x = x.view(-1, x_shape[-3], x_shape[-2], x_shape[-1]) x_embed = self.encoder(x)[0] return x_embed.view(x_shape[0], x_shape[1], -1) def reparameterize(self, mean, logvar, random_sampling=True): # Reparametrization occurs only if random sampling is set to true, otherwise mean is returned if random_sampling is True: eps = torch.randn_like(logvar) std = torch.exp(0.5 * logvar) z = mean + eps * std return z else: return mean def sample_z_prior_train(self, z_post, random_sampling=True): z_out = None z_means = None z_logvars = None batch_size = z_post.shape[0] z_t = torch.zeros(batch_size, self.z_dim).cpu() h_t_ly1 = torch.zeros(batch_size, self.hidden_dim).cpu() c_t_ly1 = torch.zeros(batch_size, self.hidden_dim).cpu() h_t_ly2 = torch.zeros(batch_size, self.hidden_dim).cpu() c_t_ly2 = torch.zeros(batch_size, self.hidden_dim).cpu() for i in range(self.frames): # two layer LSTM and two one-layer FC h_t_ly1, c_t_ly1 = self.z_prior_lstm_ly1(z_t, (h_t_ly1, c_t_ly1)) h_t_ly2, c_t_ly2 = self.z_prior_lstm_ly2(h_t_ly1, (h_t_ly2, c_t_ly2)) z_mean_t = self.z_prior_mean(h_t_ly2) z_logvar_t = self.z_prior_logvar(h_t_ly2) z_prior = self.reparameterize(z_mean_t, z_logvar_t, random_sampling) if z_out is None: # If z_out is none it means z_t is z_1, hence store it in the format [batch_size, 1, z_dim] z_out = z_prior.unsqueeze(1) z_means = z_mean_t.unsqueeze(1) z_logvars = z_logvar_t.unsqueeze(1) else: # If z_out is not none, z_t is not the initial z and hence append it to the previous z_ts collected in z_out z_out = torch.cat((z_out, z_prior.unsqueeze(1)), dim=1) z_means = torch.cat((z_means, z_mean_t.unsqueeze(1)), dim=1) z_logvars = torch.cat((z_logvars, z_logvar_t.unsqueeze(1)), dim=1) z_t = z_post[:,i,:] return z_means, z_logvars, z_out # If random sampling is true, reparametrization occurs else z_t is just set to the mean def sample_z(self, batch_size, random_sampling=True): z_out = None # This will ultimately store all z_s in the format [batch_size, frames, z_dim] z_means = None z_logvars = None # All states are initially set to 0, especially z_0 = 0 z_t = torch.zeros(batch_size, self.z_dim).cpu() # z_mean_t = torch.zeros(batch_size, self.z_dim) # z_logvar_t = torch.zeros(batch_size, self.z_dim) h_t_ly1 = torch.zeros(batch_size, self.hidden_dim).cpu() c_t_ly1 = torch.zeros(batch_size, self.hidden_dim).cpu() h_t_ly2 = torch.zeros(batch_size, self.hidden_dim).cpu() c_t_ly2 = torch.zeros(batch_size, self.hidden_dim).cpu() for _ in range(self.frames): # h_t, c_t = self.z_prior_lstm(z_t, (h_t, c_t)) # two layer LSTM and two one-layer FC h_t_ly1, c_t_ly1 = self.z_prior_lstm_ly1(z_t, (h_t_ly1, c_t_ly1)) h_t_ly2, c_t_ly2 = self.z_prior_lstm_ly2(h_t_ly1, (h_t_ly2, c_t_ly2)) z_mean_t = self.z_prior_mean(h_t_ly2) z_logvar_t = self.z_prior_logvar(h_t_ly2) z_t = self.reparameterize(z_mean_t, z_logvar_t, random_sampling) if z_out is None: # If z_out is none it means z_t is z_1, hence store it in the format [batch_size, 1, z_dim] z_out = z_t.unsqueeze(1) z_means = z_mean_t.unsqueeze(1) z_logvars = z_logvar_t.unsqueeze(1) else: # If z_out is not none, z_t is not the initial z and hence append it to the previous z_ts collected in z_out z_out = torch.cat((z_out, z_t.unsqueeze(1)), dim=1) z_means = torch.cat((z_means, z_mean_t.unsqueeze(1)), dim=1) z_logvars = torch.cat((z_logvars, z_logvar_t.unsqueeze(1)), dim=1) return z_means, z_logvars, z_out def forward(self, x, beta): _, _, f_post, _, _, z_post = self.encode_and_sample_post(x) if isinstance(f_post, list): f_expand = f_post[0].unsqueeze(1).expand(-1, self.frames, self.f_dim) else: f_expand = f_post.unsqueeze(1).expand(-1, self.frames, self.f_dim) zf = torch.cat((z_post, f_expand), dim=2) recon_x = self.decoder_frame(zf) return f_post, z_post, recon_x def name2seq(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[:, :, :3]) images = np.asarray(images, dtype='f') / 256.0 images = images.transpose((0, 3, 1, 2)) images = torch.Tensor(images).unsqueeze(dim=0) return images def display_gif(file_name, save_name): images = [] for frame in range(8): frame_name = '%d' % (frame) image_filename = file_name + frame_name + '.png' images.append(imageio.imread(image_filename)) gif_filename = 'avatar_source.gif' return imageio.mimsave(gif_filename, images) def display_gif_pad(file_name, save_name): images = [] for frame in range(8): frame_name = '%d' % (frame) image_filename = file_name + frame_name + '.png' image = imageio.imread(image_filename) image = image[:, :, :3] image_pad = cv2.copyMakeBorder(image, 0, 0, 125, 125, cv2.BORDER_CONSTANT, value=0) images.append(image_pad) return imageio.mimsave(save_name, images) def display_image(file_name): image_filename = file_name + '0' + '.png' print(image_filename) image = imageio.imread(image_filename) 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) # == Load Model == model = TransferVAE_Video(opt) model.load_state_dict(torch.load('TransferVAE.pth.tar', map_location=torch.device('cpu'))['state_dict']) model.eval() def run(domain_source, action_source, hair_source, top_source, bottom_source, domain_target, action_target, hair_target, top_target, bottom_target): # == Source Avatar == # body body_source = '0' # hair if hair_source == "green": hair_source = '0' elif hair_source == "yellow": hair_source = '2' elif hair_source == "rose": hair_source = '4' elif hair_source == "red": hair_source = '7' elif hair_source == "wine": hair_source = '8' # top if top_source == "brown": top_source = '0' elif top_source == "blue": top_source = '1' elif top_source == "white": top_source = '2' # bottom if bottom_source == "white": bottom_source = '0' elif bottom_source == "golden": bottom_source = '1' elif bottom_source == "red": bottom_source = '2' elif bottom_source == "silver": bottom_source = '3' 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) + '_' # == Target Avatar == # body body_target = '1' # hair if hair_target == "violet": hair_target = '1' elif hair_target == "silver": hair_target = '3' elif hair_target == "purple": hair_target = '5' elif hair_target == "grey": hair_target = '6' elif hair_target == "golden": hair_target = '9' # top if top_target == "grey": top_target = '3' elif top_target == "khaki": top_target = '4' elif top_target == "linen": top_target = '5' elif top_target == "ocre": top_target = '6' # bottom if bottom_target == "denim": bottom_target = '4' elif bottom_target == "olive": bottom_target = '5' elif bottom_target == "brown": bottom_target = '6' 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) + '_' # == Load Input == images_source = name2seq(file_name_source) images_target = name2seq(file_name_target) x = torch.cat((images_source, images_target), dim=0) # == Forward == with torch.no_grad(): f_post, z_post, recon_x = 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' gr.Interface( run, inputs=[ gr.Textbox(value="Source Avatar - Human", show_label=False, interactive=False), gr.Radio(choices=["slash", "spellcard", "walk"], value="slash"), gr.Radio(choices=["green", "yellow", "rose", "red", "wine"], value="green"), gr.Radio(choices=["brown", "blue", "white"], value="brown"), gr.Radio(choices=["white", "golden", "red", "silver"], value="white"), gr.Textbox(value="Target Avatar - Alien", show_label=False, interactive=False), gr.Radio(choices=["slash", "spellcard", "walk"], value="walk"), gr.Radio(choices=["violet", "silver", "purple", "grey", "golden"], value="golden"), gr.Radio(choices=["grey", "khaki", "linen", "ocre"], value="ocre"), gr.Radio(choices=["denim", "olive", "brown"], value="brown"), ], outputs=[ gr.components.Image(type="file", label="Domain Disentanglement"), ], live=True, title="TransferVAE for Unsupervised Video Domain Adaptation", ).launch()