diff --git a/models/pyramid_level.py b/models/pyramid_level.py index 62fcea1..b8abeb1 100644 --- a/models/pyramid_level.py +++ b/models/pyramid_level.py @@ -100,7 +100,7 @@ from models.upscaler import Upscaler class PyramidLevel(nn.Module): - def __init__(self, config): + def __init__(self, config, use_normals=True): super(PyramidLevel, self).__init__() self.size = config['resolution'] @@ -114,6 +114,7 @@ class PyramidLevel(nn.Module): load_model(self.upscale_enhancer, config['upscale_enhancer']) for param in self.upscale_enhancer.parameters(): param.requires_grad = False + self.use_normals = use_normals def enhance(self, low, mask): x = self.upscale_enhancer([low, mask]) @@ -152,9 +153,13 @@ class PyramidLevel(nn.Module): def forward(self, inputs, z_cond): - uv_unposed, uv_template, uv_body_posed, uv_curv, uv_mask, uv_mask_high = inputs + uv_unposed, uv_template, uv_body_posed, uv_normals, uv_mask, uv_mask_high = inputs - x = torch.concat([uv_unposed, uv_template, uv_body_posed, uv_curv], dim=1) + if self.use_normals: + x = torch.concat([uv_unposed, uv_template, uv_body_posed, uv_normals], dim=1) + else: + x = torch.concat([uv_unposed, uv_template, uv_body_posed], dim=1) + mean, logvar = self.encoder(x) z = self.reparameterize(mean, logvar) z = torch.reshape(z, (mean.shape[0], -1, 8, 4)) diff --git a/models/pyramid_model.py b/models/pyramid_model.py index da13b28..447aab5 100644 --- a/models/pyramid_model.py +++ b/models/pyramid_model.py @@ -3,11 +3,6 @@ import time import random import torch from torch import nn -import torch.nn.functional as F -# 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 @@ -15,123 +10,6 @@ from models.pyramid_level import PyramidLevel from torchvision.transforms import Resize from torchvision.transforms.functional import InterpolationMode -# class Pyramid(nn.Module): - -# def __init__(self, config): -# super(Pyramid, self).__init__() - -# self.static_enc : AEEncoder = AEEncoder(config['cond_ae_static']['encoder_block_settings'], kernel_size=config['cond_ae_static']['kernel_size'], activation=config['cond_ae_static']['activation']) -# if config['pyramid']['static_model'] != '': -# load_model(self.static_enc, config['pyramid']['static_model'] + '_enc') - -# self.body_enc : AEEncoder = AEEncoder(config['cond_ae_body']['encoder_block_settings'], kernel_size=config['cond_ae_body']['kernel_size'], activation=config['cond_ae_body']['activation']) -# if config['pyramid']['body_model'] != '': -# load_model(self.body_enc, config['pyramid']['body_model'] + '_enc') - -# self.curv_enc : AEEncoder = AEEncoder(config['cond_ae_curv']['encoder_block_settings'], kernel_size=config['cond_ae_curv']['kernel_size'], activation=config['cond_ae_curv']['activation']) -# if config['pyramid']['curvature_model'] != '': -# load_model(self.curv_enc, config['pyramid']['curvature_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 -# for param in self.curv_enc.parameters(): param.requires_grad = False - - -# self.levels = [PyramidLevel({ **config['pyramid']['levels'][i], **{ 'cond_ae_static': config['cond_ae_static'], 'cond_ae_body': config['cond_ae_body'], 'cond_ae_curv': config['cond_ae_curv'], } }) for i in range(len(config['pyramid']['levels']))] - -# # if config['is_cape']: -# # if len(config['pyramid']['resolutions']) == 2: -# # self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)]) -# # if len(config['pyramid']['resolutions']) == 3: -# # self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)]) -# # self.levels[-2].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)]) -# # if len(config['pyramid']['resolutions']) == 4: -# # self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)]) -# # self.levels[-2].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)]) -# # self.levels[-3].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)]) -# # else: - -# # if len(config['pyramid']['resolutions']) == 2: -# # self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)]) -# # if len(config['pyramid']['resolutions']) == 3: -# # self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 1000)]) -# # self.levels[-2].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)]) -# # if len(config['pyramid']['resolutions']) == 4: -# # self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 10000)]) -# # self.levels[-2].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 1000)]) -# # self.levels[-3].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)]) -# # if len(config['pyramid']['resolutions']) == 5: -# # self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100000)]) -# # self.levels[-2].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 10000)]) -# # self.levels[-3].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 1000)]) -# # self.levels[-4].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)]) - -# # self.add_noise = config['pyramid']['add_noise'] -# # self.noise_amount = config['pyramid']['noise_amount'] -# self.use_offsets = config['use_offsets'] - -# # self.upscale_enhancer = None -# # if config['pyramid']['final_upscale_enhancer'] != '': -# # k_size = config['pyramid']['final_upscale_enhancer_kernel_size'] if 'final_upscale_enhancer_kernel_size' in config['pyramid'].keys() else ((3 if config['pyramid']['resolutions'][-1] <= 64 else 5)) -# # self.upscale_enhancer = UpscalerBasicConv(low_res=config['pyramid']['resolutions'][-1], kernel_size=k_size) -# # load_model(self.upscale_enhancer, config['pyramid']['final_upscale_enhancer']) -# # self.upscale_enhancer.trainable = False - -# def enhance(self, low, mask): -# x = self.upscale_enhancer([low, mask]) -# return x - -# def encode_conds(self, uv_template, uv_body_posed, uv_curv): -# z_static = self.static_enc(uv_template) -# z_body = self.body_enc(uv_body_posed) -# z_curv = self.curv_enc(uv_curv) - -# z_cond = torch.concat([z_static, z_body, z_curv], dim=-1) - -# return z_cond - -# def call(self, inputs): -# levels_input, conds = inputs -# uv_static, uv_body, uv_curv = conds[0], conds[1], conds[2] - -# z_cond = self.encode_conds(uv_static, uv_body, uv_curv) - -# # if self.add_noise == True and training == True: -# # z_cond = z_cond + (tf.random.uniform(z_cond.shape, maxval=self.noise_amount * 2) - self.noise_amount) - -# out_levels = [] -# pred, mean, logvar = self.levels[0](levels_input[0], z_cond) -# if self.use_offsets: -# pred = pred + levels_input[0][2] # uv_body_posed_32 -# if 1 != len(self.levels): -# pred = self.levels[0].enhance(pred, levels_input[0][-2]) # levels_input[i][-2]=uv_mask_res -# pred = F.interpolate(pred, size=(self.levels[1].size, self.levels[1].size//2), mode='bilinear') -# out_levels.append([pred, mean, logvar]) -# pred = pred * levels_input[1][-2][:,None,:,:] -# else: -# pred = self.enhance(pred, levels_input[0][-2]) -# pred = F.interpolate(pred, size=(512, 256), mode='bilinear') -# pred = pred * levels_input[-1][-1][:,None,:,:] # full res uv mask -# out_levels.append([pred, mean, logvar]) - -# for i in range(1, len(self.levels)): -# pred_level, mean, logvar = self.levels[i](levels_input[i], z_cond) -# pred = pred + pred_level -# if i != len(self.levels) - 1: -# pred = self.levels[i].enhance(pred, levels_input[i][-2]) -# out_levels.append([pred, mean, logvar]) -# pred = F.interpolate(pred, size=(self.levels[i+1].size, self.levels[i+1].size//2), mode='bilinear') -# pred = pred * levels_input[i+1][-2][:,None,:,:] -# else: -# pred = self.enhance(pred, levels_input[i][-1]) -# pred = pred * levels_input[-1][-1][:,None,:,:] # full res uv mask -# pred = F.interpolate(pred, size=(512, 256), mode='bilinear') -# out_levels.append([pred, mean, logvar]) - -# return out_levels - - class Pyramid(nn.Module): @@ -146,57 +24,33 @@ class Pyramid(nn.Module): if config['body_model'] != '': load_model(self.body_enc, config['body_model'] + '_enc') - self.norm_enc : AEEncoder = AEEncoder(config['cond_ae_norm']['encoder_block_settings'], config['cond_ae_norm']['kernel_size'], config['cond_ae_norm']['activation'], device=config['device']) - if config['normals_model'] != '': - load_model(self.norm_enc, config['normals_model'] + '_enc') + if config['use_normals']: + self.norm_enc : AEEncoder = AEEncoder(config['cond_ae_norm']['encoder_block_settings'], config['cond_ae_norm']['kernel_size'], config['cond_ae_norm']['activation'], device=config['device']) + if config['normals_model'] != '': + load_model(self.norm_enc, config['normals_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 - for param in self.norm_enc.parameters(): param.requires_grad = False + if config['use_normals']: + for param in self.norm_enc.parameters(): param.requires_grad = False - self.levels = nn.ModuleList([PyramidLevel({ **config['pyramid']['levels'][i] }) for i in range(len(config['pyramid']['levels']))]) + self.levels = nn.ModuleList([PyramidLevel({ **config['pyramid']['levels'][i] }, use_normals=config['use_normals']) for i in range(len(config['pyramid']['levels']))]) self.resizers = { res: Resize(size=(res, int(res/2)), interpolation=InterpolationMode.BILINEAR).to(config['device']) for res in config['pyramid']['resolutions'] } self.resizers[512] = Resize(size=(512, 256), interpolation=InterpolationMode.BILINEAR).to(config['device']) - - # if len(config['pyramid']['resolutions']) == 2: - # self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)]) - # if len(config['pyramid']['resolutions']) == 3: - # self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 1000)]) - # self.levels[-2].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)]) - # if len(config['pyramid']['resolutions']) == 4: - # self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 10000)]) - # self.levels[-2].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 1000)]) - # self.levels[-3].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)]) - # if len(config['pyramid']['resolutions']) == 5: - # self.levels[-1].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100000)]) - # self.levels[-2].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 10000)]) - # self.levels[-3].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 1000)]) - # self.levels[-4].decoder.decoder.last = tf.keras.Sequential([tf.keras.layers.Dense(3), tf.keras.layers.Lambda(lambda x: x / 100)]) - - # self.add_noise = config['pyramid']['add_noise'] - # self.noise_amount = config['pyramid']['noise_amount'] - # self.use_offsets = config['use_offsets'] - - # self.upscale_enhancer = None - # if config['pyramid']['final_upscale_enhancer'] != '': - # k_size = config['pyramid']['final_upscale_enhancer_kernel_size'] if 'final_upscale_enhancer_kernel_size' in config['pyramid'].keys() else ((3 if config['pyramid']['resolutions'][-1] <= 64 else 5)) - # self.upscale_enhancer = UpscalerBasicConv(low_res=config['pyramid']['resolutions'][-1], kernel_size=k_size) - # load_model(self.upscale_enhancer, config['pyramid']['final_upscale_enhancer']) - # self.upscale_enhancer.trainable = False - - # def enhance(self, low, mask, training): - # x = self.upscale_enhancer([low, mask], training=training) - # return x + self.use_normals = config['use_normals'] def encode_conds(self, uv_template, uv_body_posed, uv_norm): z_static = self.static_enc(uv_template) z_body = self.body_enc(uv_body_posed) - z_norm = self.norm_enc(uv_norm) - z_cond = torch.concat([z_static, z_body, z_norm], dim=1) + if self.use_normals: + z_norm = self.norm_enc(uv_norm) + z_cond = torch.concat([z_static, z_body, z_norm], dim=1) + else: + z_cond = torch.concat([z_static, z_body], dim=1) return z_cond @@ -206,9 +60,6 @@ class Pyramid(nn.Module): z_cond = self.encode_conds(uv_static, uv_body, uv_norm) - # if self.add_noise == True and training == True: - # z_cond = z_cond + (tf.random.uniform(z_cond.shape, maxval=self.noise_amount * 2) - self.noise_amount) - out_levels = [] pred, mean, logvar = self.levels[0](levels_input[0], z_cond) # if self.use_offsets: @@ -219,7 +70,6 @@ class Pyramid(nn.Module): out_levels.append([pred, mean, logvar]) pred = pred * levels_input[1][-2][:,None,:,:] else: - # pred = self.enhance(pred, levels_input[0][-2], training=training) pred = self.levels[0].enhance(pred, levels_input[0][-2]) # levels_input[i][-2]=uv_mask_res pred = self.resizers[512](pred) pred = pred * levels_input[-1][-1][:,None,:,:] # full res uv mask @@ -227,12 +77,8 @@ class Pyramid(nn.Module): for i in range(1, len(self.levels)): input = levels_input[i] - # input[0] = input[0] - pred - # input_min, input_max = tf.reduce_min(input[0], axis=[1,2,3]), tf.reduce_max(input[0], axis=[1,2,3]) - # input[0] = (input[0] - input_min[:,None,None,None]) / (input_max - input_min)[:,None,None,None] - # input[0] = input[0] * input[-2][:,None,:,:] # uv_mask_res + input[0] = input[0] - pred pred_level, mean, logvar = self.levels[i](input, z_cond) - # pred_level = ( (pred_level * (input_max - input_min)[:,None,None,None]) + input_min[:,None,None,None] ) * input[-2][:,:,:,None] # uv_mask_res pred = pred + pred_level if i != len(self.levels) - 1: pred = self.levels[i].enhance(pred, levels_input[i][-2]) @@ -240,8 +86,6 @@ class Pyramid(nn.Module): pred = self.resizers[self.levels[i+1].size](pred) pred = pred * levels_input[i+1][-2][:,None,:,:] else: - # if self.upscale_enhancer is not None: - # pred = self.enhance(pred, levels_input[-1][-2], training=training) pred = self.levels[-1].enhance(pred, levels_input[-1][-2]) pred = self.resizers[512](pred) pred = pred * levels_input[-1][-1][:,None,:,:] # full res uv mask diff --git a/scripts/cvae_pyramid.py b/scripts/cvae_pyramid.py index 48fbdc3..9779384 100644 --- a/scripts/cvae_pyramid.py +++ b/scripts/cvae_pyramid.py @@ -11,7 +11,7 @@ parser.add_argument('--debug', action='store_true') args = vars(parser.parse_args()) if args['debug']: - from configs.iccvw.pyr_64_cape_m import config + from configs.iccvw.pyr_64_reg import config else: config_module = importlib.import_module(args['config']) config = config_module.config @@ -238,72 +238,41 @@ def train_step(batch): res_modifier = '_' + str(res) if res != 512 else '' pred, mean, logvar = out_levels[i] - mesh_pred = uv2mesh(pred) - if config['use_skinning']: - shape = batch['pose_shape'][:, 72:] - gender = torch.greater_equal(shape[:, 10] , 0.5).to(dtype=torch.int64) - shape_offsets = torch.stack([torch.tensordot(shape[i,:10][None, :], smpl.shapedirs[gender[i]], dims=[[1], [2]])[0,smpl.to_keep,:] for i in range(shape.shape[0])]) - mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,None] - mesh_pred_posed = smpl(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)[0] - else: - mesh_pred_posed = mesh_pred - - # is the output already masked? - # ms_ssim = torch.mean(tf.image.ssim_multiscale(pred + 5, batch['uv_unposed'] + 5, max_val=10, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1))) - # ms_ssim = torch.mean(tf.image.ssim(pred + 1, batch['uv_unposed'] + 1, max_val=3)) - - e_loss = torch.mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask'])) - c_loss, c_ratio = collision_loss(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask']) - c_loss, c_ratio = torch.mean(c_loss), torch.mean(c_ratio) l2_loss = torch.mean(1e3 * torch.mean(torch.sum(((uv_orig - pred) ** 2), dim=-1), dim=[1, 2])) - l2_mesh_loss = torch.mean(1e3 * torch.mean(torch.sum(batch['mesh_mask'][:,:,None] * (batch['mesh_posed'] - mesh_pred_posed) ** 2, dim=-1), dim=-1)) - n_loss = normal_loss(mesh_pred_posed, batch['mesh_posed'], batch['mesh_mask'], is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape']) kl = kl_loss(mean, logvar) if config['use_extra_losses']: + mesh_pred = uv2mesh(pred) + + if config['use_skinning']: + shape = batch['pose_shape'][:, 72:] + gender = torch.greater_equal(shape[:, 10] , 0.5).to(dtype=torch.int64) + shape_offsets = torch.stack([torch.tensordot(shape[i,:10][None, :], smpl.shapedirs[gender[i]], dims=[[1], [2]])[0,smpl.to_keep,:] for i in range(shape.shape[0])]) + + mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,None] + mesh_pred_posed = smpl(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)[0] + else: + mesh_pred_posed = mesh_pred + + e_loss = torch.mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask'])) + c_loss, c_ratio = collision_loss(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask']) + c_loss, c_ratio = torch.mean(c_loss), torch.mean(c_ratio) + n_loss = normal_loss(mesh_pred_posed, batch['mesh_posed'], batch['mesh_mask'], is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape']) + l2_mesh_loss = torch.mean(1e3 * torch.mean(torch.sum(batch['mesh_mask'][:,:,None] * (batch['mesh_posed'] - mesh_pred_posed) ** 2, dim=-1), dim=-1)) + reconstr_loss_total = reconstr_loss_total + config['pyramid']['loss_weights'][i] * (l2_loss + l2_mesh_loss) if not config['is_cape']: reconstr_loss_total = reconstr_loss_total + config['pyramid']['loss_weights'][i] * (e_loss + c_loss) * config['reg_loss_coeff'] + if config['use_norm_loss'] and step > 2000: + reconstr_loss_total = reconstr_loss_total + n_loss * config['reg_loss_coeff'] else: reconstr_loss_total = reconstr_loss_total + config['pyramid']['loss_weights'][i] * (l2_loss) - - if config['use_norm_loss'] and step > 2000: - reconstr_loss_total = reconstr_loss_total + n_loss * config['reg_loss_coeff'] kl_total += kl loss: torch.Tensor = reconstr_loss_total + g_balancing_coeff * kl_total - # scaler.scale(loss).backward() - # scaler.step(optimizer) - # scaler.update() - # optimizer.zero_grad() - - # b = 0 - # from custom_utils.custom_utils import _maskFaces - # mesh_mask = batch['mesh_mask'][b].cpu().numpy() - # faces = np.load('smpl/smpl_faces_skirt.npy') if batch['garment'][b] in ['Skirt', 'Dress'] else np.load('smpl/smpl_faces.npy') - # faces_body = np.load('smpl/smpl_faces.npy') - # _, faces = _maskFaces(np.zeros((14475, 3)), faces, mesh_mask) - - # # fig = go.Figure(data=[go.Mesh3d(x=batch['mesh_posed'].cpu().numpy()[0,mesh_mask,0], y=batch['mesh_posed'].cpu().numpy()[0,mesh_mask,2], z=batch['mesh_posed'].cpu().numpy()[0,mesh_mask,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightpink', opacity=1), - # # go.Mesh3d(x=batch['mesh_body_posed'].cpu().numpy()[0,:,0], y=batch['mesh_body_posed'].cpu().numpy()[0,:,2], z=batch['mesh_body_posed'].cpu().numpy()[0,:,1], i=faces_body[:,0], j=faces_body[:,1], k=faces_body[:,2], color='lightblue', opacity=1), - # # ]) - - # fig = go.Figure(data=[go.Mesh3d(x=mesh_orig.cpu().numpy()[0,mesh_mask,0], y=mesh_orig.cpu().numpy()[0,mesh_mask,2], z=mesh_orig.cpu().numpy()[0,mesh_mask,1], i=faces[:,0], j=faces[:,1], k=faces[:,2], color='lightpink', opacity=1), - # ]) - # fig.update_layout(scene_aspectmode='cube') - # fig.update_layout(scene=dict(xaxis=dict(range=[-1,1]), yaxis=dict(range=[-1,1]), zaxis=dict(range=[-1,1]))) - # fig.show() - - # ex.imshow((uv_orig[b].cpu().numpy().transpose(1,2,0) + 1) / 2).show() - # upscaled = model.levels[0].enhance(levels_input[0][0], levels_input[0][-2]) # levels_input[i][-2]=uv_mask_res - # upscaled = model.resizers[512](upscaled) - # ex.imshow((levels_input[0][0][b].cpu().numpy().transpose(1,2,0) + 1) / 2).show() - # ex.imshow((upscaled[b].cpu().numpy().transpose(1,2,0) + 1) / 2).show() - - optimizer.zero_grad() loss.backward() optimizer.step() @@ -313,13 +282,33 @@ def train_step(batch): mean = torch.mean(mean) logvar = torch.mean(logvar) + if not config['use_extra_losses']: + mesh_pred = uv2mesh(pred) + + if config['use_skinning']: + shape = batch['pose_shape'][:, 72:] + gender = torch.greater_equal(shape[:, 10] , 0.5).to(dtype=torch.int64) + shape_offsets = torch.stack([torch.tensordot(shape[i,:10][None, :], smpl.shapedirs[gender[i]], dims=[[1], [2]])[0,smpl.to_keep,:] for i in range(shape.shape[0])]) + + mesh_pred_shaped = mesh_pred + shape_offsets * batch['mesh_mask'][:,:,None] + mesh_pred_posed = smpl(batch['pose_shape'][:, :72], batch['pose_shape'][:, 72:], mesh_pred_shaped)[0] + else: + mesh_pred_posed = mesh_pred + + e_loss = torch.mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask'])) + c_loss, c_ratio = collision_loss(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask']) + c_loss, c_ratio = torch.mean(c_loss), torch.mean(c_ratio) + l2_mesh_loss = torch.mean(1e3 * torch.mean(torch.sum(batch['mesh_mask'][:,:,None] * (batch['mesh_posed'] - mesh_pred_posed) ** 2, dim=-1), dim=-1)) + n_loss = normal_loss(mesh_pred_posed, batch['mesh_posed'], batch['mesh_mask'], is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape']) + ms_ssim = torch.mean(_MSSSIM(pred + 1, batch['uv_static'] + 1)) metric = torch.sum(torch.sqrt(torch.sum(batch['mesh_mask'][:,:,None] * (mesh_orig - mesh_pred.detach()) ** 2, axis=-1))) / torch.sum(batch['mesh_mask']) output = { **output, **{ 'train/edge_loss' + res_modifier: e_loss, 'train/collision_loss' + res_modifier: c_loss,'train/euclid_error' + res_modifier: metric, 'train/mean' + res_modifier: mean, 'train/logvar' + res_modifier: logvar, 'train/kl_loss' + res_modifier: kl, 'train/var' + res_modifier: np.e ** logvar, 'train/l2_loss' + res_modifier: l2_loss, 'train/normal_loss': n_loss, - 'train/euclid_error': metric, 'train/l2_mesh_loss': l2_mesh_loss, 'train/ms_ssim' + res_modifier: ms_ssim, } } # + 'train/euclid_error': metric, 'train/l2_mesh_loss': l2_mesh_loss, 'train/ms_ssim' + res_modifier: ms_ssim, + 'train/collision_ratio': c_ratio } } # return { **{ 'train/loss': loss, 'train/reconstr_loss_total': reconstr_loss_total, 'train/kl_total': kl_total }, **output } @@ -359,28 +348,24 @@ def valid_step(batch): else: mesh_pred_posed = mesh_pred - # is the output already masked? - # ms_ssim = torch.mean(tf.image.ssim_multiscale(pred + 5, batch['uv_unposed'] + 5, max_val=10, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1))) - # ms_ssim = torch.mean(tf.image.ssim(pred + 1, batch['uv_unposed'] + 1, max_val=3)) - e_loss = torch.mean(edge_loss(mesh_pred, mesh_orig, batch['mesh_mask'])) c_loss, c_ratio = collision_loss(mesh_pred_posed, batch['mesh_body_posed'], batch['mesh_mask']) c_loss, c_ratio = torch.mean(c_loss), torch.mean(c_ratio) l2_loss = torch.mean(1e3 * torch.mean(torch.sum(((uv_orig - pred) ** 2), dim=-1), dim=[1, 2])) l2_mesh_loss = torch.mean(1e3 * torch.mean(torch.sum(batch['mesh_mask'][:,:,None] * (batch['mesh_posed'] - mesh_pred_posed) ** 2, dim=-1), dim=-1)) n_loss = normal_loss(mesh_pred_posed, batch['mesh_posed'], batch['mesh_mask'], is_skirt=config['is_skirt'], is_smpl_orig=config['is_cape']) + kl = kl_loss(mean, logvar) if config['use_extra_losses']: reconstr_loss_total = reconstr_loss_total + config['pyramid']['loss_weights'][i] * (l2_loss + l2_mesh_loss) if not config['is_cape']: reconstr_loss_total = reconstr_loss_total + config['pyramid']['loss_weights'][i] * (e_loss + c_loss) * config['reg_loss_coeff'] + if config['use_norm_loss'] and step > 2000: + reconstr_loss_total = reconstr_loss_total + n_loss * config['reg_loss_coeff'] else: reconstr_loss_total = reconstr_loss_total + config['pyramid']['loss_weights'][i] * (l2_loss) - if config['use_norm_loss'] and step > 2000: - reconstr_loss_total = reconstr_loss_total + n_loss * config['reg_loss_coeff'] - kl_total += kl loss = reconstr_loss_total + g_balancing_coeff * kl_total @@ -825,9 +810,9 @@ def train(): print(step) if train_loss['train/euclid_error'] < config['balancing_activation_limit']: - g_balancing_coeff += 1 / 10000 + g_balancing_coeff += 1 / 50000 - if step % 25000 == 0: + if step % 25000 == 0 and not args['debug']: ckpt_saver(model, step, optimizer, None) step += 1 @@ -874,12 +859,14 @@ if __name__ == '__main__': if len(config['pyramid']['resolutions']) == 1: for param in model.static_enc.parameters(): param.requires_grad = False for param in model.body_enc.parameters(): param.requires_grad = False - for param in model.norm_enc.parameters(): param.requires_grad = False + if config['use_normals']: + for param in model.norm_enc.parameters(): param.requires_grad = False else: - load_model(model, config['pyramid']['pretrained_pyramid'], iteration=(config['pyramid']['pretrained_pyramid_iteration'] if 'pretrained_pyramid_iteration' in config.keys() else '')) + load_model(model, config['pyramid']['pretrained_pyramid'], iteration=config['pyramid']['pretrained_pyramid_iteration']) for param in model.static_enc.parameters(): param.requires_grad = False for param in model.body_enc.parameters(): param.requires_grad = False - for param in model.norm_enc.parameters(): param.requires_grad = False + if config['use_normals']: + for param in model.norm_enc.parameters(): param.requires_grad = False # load_model(model.levels[0].upscale_enhancer, config['pyramid']['levels'][0]['upscale_enhancer']) # model.levels[0].upscale_enhancer.trainable = False