|
import gradio as gr |
|
|
|
import cv2 |
|
import imageio |
|
import math |
|
from math import ceil |
|
import matplotlib.pyplot as plt |
|
import numpy as np |
|
from PIL import Image |
|
import subprocess |
|
import torch |
|
import torch.nn as nn |
|
import torch.nn.functional as F |
|
|
|
|
|
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 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 encode_and_sample_post(self, x): |
|
conv_x = self.encoder_frame(x) |
|
lstm_out, _ = self.z_lstm(conv_x) |
|
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_post = f_mean |
|
|
|
features, _ = self.z_rnn(lstm_out) |
|
|
|
z_mean = self.z_mean(features) |
|
z_post = z_mean |
|
|
|
return f_post, 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 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 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.0125, wspace=0.0) |
|
|
|
save_name = 'MyPlot_{}.png'.format(frame_id) |
|
|
|
plt.savefig(save_name, dpi=200, format='png', bbox_inches='tight', pad_inches=0.0) |
|
|
|
|
|
|
|
model = TransferVAE_Video() |
|
model.load_state_dict(torch.load('TransferVAE.pth.tar', map_location=torch.device('cpu'))['state_dict']) |
|
model.eval() |
|
|
|
|
|
def run(source, action_source, hair_source, top_source, bottom_source, target, action_target, hair_target, top_target, bottom_target): |
|
|
|
|
|
|
|
body_source = '0' |
|
|
|
|
|
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' |
|
|
|
|
|
if top_source == "brown": top_source = '0' |
|
elif top_source == "blue": top_source = '1' |
|
elif top_source == "white": top_source = '2' |
|
|
|
|
|
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) + '_' |
|
|
|
|
|
|
|
|
|
body_target = '1' |
|
|
|
|
|
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' |
|
|
|
|
|
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' |
|
|
|
|
|
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) + '_' |
|
|
|
|
|
|
|
images_source = name2seq(file_name_source) |
|
images_target = name2seq(file_name_target) |
|
x = torch.cat((images_source, images_target), dim=0) |
|
|
|
|
|
|
|
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) |
|
|
|
|
|
|
|
for frame in range(8): |
|
|
|
|
|
src_orig = src_orig_sample[frame, :, :, :].detach().numpy().transpose((1, 2, 0)) |
|
tar_orig = tar_orig_sample[frame, :, :, :].detach().numpy().transpose((1, 2, 0)) |
|
|
|
|
|
src_recon = src_recon_sample[frame, :, :, :].detach().numpy().transpose((1, 2, 0)) |
|
tar_recon = tar_recon_sample[frame, :, :, :].detach().numpy().transpose((1, 2, 0)) |
|
|
|
|
|
f_expand_src = 0 * src_f_post.unsqueeze(1).expand(-1, 8, 512) |
|
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, 512) |
|
zf_tar = torch.cat((tar_z_post, f_expand_tar), dim=2) |
|
recon_x_tar = model.decoder_frame(zf_tar) |
|
tar_Zt = recon_x_tar.squeeze()[frame, :, :, :].detach().numpy().transpose((1, 2, 0)) |
|
|
|
|
|
f_expand_src = src_f_post.unsqueeze(1).expand(-1, 8, 512) |
|
zf_srcZf_tarZt = torch.cat((tar_z_post, f_expand_src), dim=2) |
|
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, 512) |
|
zf_tarZf_srcZt = torch.cat((src_z_post, f_expand_tar), dim=2) |
|
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' |
|
|
|
|
|
desc = """ |
|
Welcome to the demo page of TranSVAE, a disentanglement framework designed for unsupervised video domain adaptation. In this live demo, you are able to: |
|
|
|
- Explore domain disentanglement and transfer in TranSVAE with Sprites avatars; |
|
- Customize the Sprites avatars by yourself via changing their actions, hair colors, top wears, and bottom wears. |
|
|
|
For more details, read the [TranSVAE paper](https://arxiv.org/abs/2208.07365) and visit our [project page](https://ldkong.com/TranSVAE). The training and testing code is available at our [GitHub Repo](https://github.com/ldkong1205/TranSVAE). Have fun! |
|
""" |
|
|
|
gr.Interface( |
|
fn=run, |
|
inputs=[ |
|
gr.Markdown( |
|
""" |
|
👦🏻 Human - Source Avatar |
|
""" |
|
), |
|
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.Markdown( |
|
""" |
|
👽 Alien - Target Avatar |
|
""" |
|
), |
|
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=False, |
|
cache_examples=True, |
|
title="TranSVAE for Unsupervised Video Domain Adaptation", |
|
description=desc |
|
).launch(share=True) |
|
|
