97 lines
3.7 KiB
Python
97 lines
3.7 KiB
Python
import os
|
|
import time
|
|
import random
|
|
import torch
|
|
from torch import nn
|
|
import numpy as np
|
|
import scipy as sp
|
|
# import plotly.graph_objects as go
|
|
# import plotly.express as ex
|
|
|
|
from utils.checkpoints import load_model
|
|
from models.base import AEEncoder, VariationalEncoder, Decoder
|
|
|
|
|
|
class CVAENorm(nn.Module):
|
|
|
|
def __init__(self, config):
|
|
super(CVAENorm, self).__init__()
|
|
self.size = config['resolution']
|
|
|
|
self.static_enc : AEEncoder = AEEncoder(config['cond_ae_static']['encoder_block_settings'], config['cond_ae_static']['kernel_size'], config['cond_ae_static']['activation'], device=config['device'])
|
|
if config['static_model'] != '':
|
|
load_model(self.static_enc, config['static_model'] + '_enc')
|
|
|
|
self.body_enc : AEEncoder = AEEncoder(config['cond_ae_body']['encoder_block_settings'], config['cond_ae_body']['kernel_size'], config['cond_ae_body']['activation'], device=config['device'])
|
|
if config['body_model'] != '':
|
|
load_model(self.body_enc, config['body_model'] + '_enc')
|
|
|
|
if config['freeze_encoders']:
|
|
for param in self.static_enc.parameters(): param.requires_grad = False
|
|
for param in self.body_enc.parameters(): param.requires_grad = False
|
|
|
|
self.encoder = VariationalEncoder(config['cvae']['encoder_block_settings'], config['cvae']['kernel_size'], config['cvae']['input_features'], device=config['device'])
|
|
self.decoder : Decoder = Decoder(config['cvae']['decoder_block_settings'], kernel_size=config['cvae']['kernel_size'], device=config['device'])
|
|
|
|
self.latent_shape = (config['cvae']['latent_dim'], 8, 4)
|
|
|
|
|
|
def sample(self, uv_template, uv_body_posed, uv_curv, uv_mask, noise=None, add_noise=False):
|
|
if noise is None:
|
|
eps = tf.random.normal(shape=(uv_template.shape[0], 8, 4, self.latent_dim))
|
|
else:
|
|
eps = noise
|
|
|
|
z_cond = self.encode_conds(uv_template, uv_body_posed, uv_curv, training=False)
|
|
if add_noise:
|
|
z_cond = z_cond + (tf.random.uniform(z_cond.shape, maxval=self.noise_amount * 2) - self.noise_amount)
|
|
|
|
z = tf.concat([eps, z_cond], axis=-1)
|
|
|
|
x = self.decoder(z, training=False)
|
|
|
|
if self.avg_garment is not None:
|
|
x = x + self.avg_garment[tf.newaxis,:,:,:]
|
|
|
|
if self.use_offsets_over_body:
|
|
x = x + uv_body_posed
|
|
|
|
x = x * uv_mask[:,:,:,tf.newaxis]
|
|
|
|
return x
|
|
|
|
|
|
def reparameterize(self, mean: torch.Tensor, logvar: torch.Tensor, training: bool = True, fix_noise: bool = False):
|
|
if training == True:
|
|
if not fix_noise:
|
|
eps = torch.randn(size=mean.shape).to(mean.device)
|
|
else:
|
|
eps = torch.randn(size=(1, mean.shape[1]))
|
|
eps = torch.tile(eps, [mean.shape[0],1])
|
|
out = eps * torch.exp(logvar * 0.5) + mean
|
|
else:
|
|
out = torch.exp(logvar * 0.5) + mean
|
|
return out
|
|
|
|
|
|
def encode_conds(self, uv_template, uv_body_posed):
|
|
z_body = self.body_enc(uv_body_posed)
|
|
z_static = self.static_enc(uv_template)
|
|
z_cond = torch.concat([z_static, z_body], dim=1)
|
|
return z_cond
|
|
|
|
|
|
def forward(self, inputs):
|
|
uv_unposed, uv_template, uv_body_posed, uv_mask = inputs
|
|
x = torch.concat([uv_unposed, uv_template, uv_body_posed], dim=1)
|
|
|
|
mean, logvar = self.encoder(x)
|
|
z: torch.Tensor = self.reparameterize(mean, logvar)
|
|
z = z.view((mean.shape[0], *self.latent_shape))
|
|
z_cond = self.encode_conds(uv_template, uv_body_posed)
|
|
|
|
z = torch.concat([z, z_cond], dim=1)
|
|
x = self.decoder(z)
|
|
x = x * uv_mask[:,None,:,:]
|
|
|
|
return x, mean, logvar |