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

small updates

Browse Source
This commit is contained in:
luni
2023-07-08 11:54:21 +02:00
parent c41a577672
commit 57b7aed4ba
15 changed files with 884 additions and 221 deletions
Download Patch File Download Diff File Expand all files Collapse all files
models/pyramid_level.py
+3 -85
View File
@@ -13,91 +13,6 @@ from models.base import AEEncoder, VariationalEncoder, Decoder
from models.upscaler import Upscaler
# class PyramidLevel(nn.Module):
# def __init__(self, config):
# super(PyramidLevel, self).__init__()
# self.latent_dim = config['latent_dim']
# self.size = config['resolution']
# self.encoder: VariationalEncoder = VariationalEncoder(config['encoder_block_settings'], config['input_features'], kernel_size=config['kernel_size'])
# self.decoder: Decoder = Decoder(config['decoder_block_settings'], config['kernel_size'])
# self.upscale_enhancer = None
# if config['upscale_enhancer'] != '':
# k_size = config['upscale_enhancer_kernel_size']
# self.upscale_enhancer = Upscaler(config['resolution'], kernel_size=k_size)
# load_model(self.upscale_enhancer, config['upscale_enhancer'])
# for param in self.upscale_enhancer.parameters(): param.requires_grad = False
# if config['avg_garment_path'] != '':
# self.avg_garment = np.load(config['avg_garment_path'])
# else:
# self.avg_garment = None
# def enhance(self, low, mask, training=None):
# x = self.upscale_enhancer([low, mask], training=training)
# return x
# def sample(self, z_cond, uv_mask, noise=None):
# if noise is None:
# eps = torch.normal(shape=(uv_mask.shape[0], self.latent_dim))
# else:
# eps = noise
# z = torch.concat([eps, z_cond], dim=-1)
# x = self.decoder(z)
# if self.avg_garment is not None:
# x = x + self.avg_garment[None,:,:,:]
# x = x * uv_mask[:,None,:,:]
# return x
# def reparameterize(self, mean, logvar, training, fix_noise=False):
# if training == True:
# if not fix_noise:
# eps = torch.normal(mean=0.0, std=1.0, size=mean.shape)
# else:
# eps = torch.normal(mean=0.0, std=1.0, size=mean.shape)
# eps = torch.tile(eps, [mean.shape[0],1])
# out = eps * torch.exp(logvar * .5) + mean
# else:
# out = torch.exp(logvar * .5) + mean
# return out
# def call(self, inputs, z_cond, training=None):
# uv_unposed, uv_template, uv_body_posed, uv_curv, uv_mask, uv_mask_high = inputs
# x = torch.concat([uv_unposed, uv_template, uv_body_posed, uv_curv], dim=1)
# mean, logvar = self.encoder(x, training=training)
# z = self.reparameterize(mean, logvar, training)
# z = torch.concat([z, z_cond], dim=-1)
# x = self.decoder(z, training=training)
# if self.avg_garment is not None:
# x = x + self.avg_garment[None,:,:,:]
# x = x * uv_mask[:,None,:,:]
# # if self.upscale_enhancer is not None:
# # pred_low = self.enhance(x, uv_mask, training=training)
# # x = tf.image.resize(pred_low, (512, 256), tf.image.ResizeMethod.LANCZOS5)
# # x = x * uv_mask_high[:,:,:,tf.newaxis]
# return x, mean, logvar
class PyramidLevel(nn.Module):
def __init__(self, config, use_normals=True):
@@ -115,6 +30,7 @@ class PyramidLevel(nn.Module):
for param in self.upscale_enhancer.parameters(): param.requires_grad = False
self.use_normals = use_normals
self.scaling = 1.0 if 'scaling' not in config else config['scaling']
def enhance(self, low, mask):
x = self.upscale_enhancer([low, mask])
@@ -175,4 +91,6 @@ class PyramidLevel(nn.Module):
# x = x * uv_mask_high[:,:,:,tf.newaxis]
x = x * self.scaling
return x, mean, logvar
models/pyramid_model.py
+4 -3
View File
@@ -66,8 +66,8 @@ 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 = self.resizers[self.levels[1].size](pred)
out_levels.append([pred, mean, logvar])
pred = self.resizers[self.levels[1].size](pred)
pred = pred * levels_input[1][-2][:,None,:,:]
else:
pred = self.levels[0].enhance(pred, levels_input[0][-2]) # levels_input[i][-2]=uv_mask_res
@@ -86,8 +86,9 @@ class Pyramid(nn.Module):
pred = self.resizers[self.levels[i+1].size](pred)
pred = pred * levels_input[i+1][-2][:,None,:,:]
else:
pred = self.levels[-1].enhance(pred, levels_input[-1][-2])
pred = self.resizers[512](pred)
if self.levels[-1].size != 512:
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
out_levels.append([pred, mean, logvar])