luni/cvpr_2023
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

pyramid training working

Browse Source
This commit is contained in:
luni
2023-07-02 23:23:41 +02:00
parent 2fc0177d5a
commit 592b8df0f0
4 changed files with 668 additions and 30 deletions
Download Patch File Download Diff File Expand all files Collapse all files
models/pyramid_model.py
+18 -14
View File
@@ -12,7 +12,8 @@ import torch.nn.functional as F
from utils.checkpoints import load_model
from models.base import AEEncoder
from models.pyramid_level import PyramidLevel
from torchvision.transforms import Resize
from torchvision.transforms.functional import InterpolationMode
# class Pyramid(nn.Module):
@@ -145,17 +146,20 @@ class Pyramid(nn.Module):
if config['body_model'] != '':
load_model(self.body_enc, config['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'], is_label=True, nr_of_labels=7)
if config['curvature_model'] != '':
load_model(self.curv_enc, config['curvature_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['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
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.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)])
@@ -187,20 +191,20 @@ class Pyramid(nn.Module):
# x = self.upscale_enhancer([low, mask], training=training)
# return x
def encode_conds(self, uv_template, uv_body_posed, uv_curv):
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_curv = self.curv_enc(uv_curv)
z_norm = self.norm_enc(uv_norm)
z_cond = torch.concat([z_static, z_body, z_curv], dim=1)
z_cond = torch.concat([z_static, z_body, z_norm], dim=1)
return z_cond
def forward(self, inputs):
levels_input, conds = inputs
uv_static, uv_body, uv_curv = conds[0], conds[1], conds[2]
uv_static, uv_body, uv_norm = conds[0], conds[1], conds[2]
z_cond = self.encode_conds(uv_static, uv_body, uv_curv)
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)
@@ -211,13 +215,13 @@ class Pyramid(nn.Module):
# 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')
pred = self.resizers[self.levels[1].size](pred)
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 = F.interpolate(pred, size=(512, 256), mode='bilinear')
pred = self.resizers[512](pred)
pred = pred * levels_input[-1][-1][:,None,:,:] # full res uv mask
out_levels.append([pred, mean, logvar])
@@ -233,13 +237,13 @@ class Pyramid(nn.Module):
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 = 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 = F.interpolate(pred, size=(512, 256), mode='bilinear')
pred = self.resizers[512](pred)
pred = pred * levels_input[-1][-1][:,None,:,:] # full res uv mask
out_levels.append([pred, mean, logvar])