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

save all

Browse Source
This commit is contained in:
luni
2023-08-24 11:36:22 +00:00
parent 59f67a2495
commit afd482522d
36 changed files with 2194 additions and 291 deletions
Download Patch File Download Diff File Expand all files Collapse all files
models/pyramid_level.py
+1 -1
View File
@@ -39,7 +39,7 @@ class PyramidLevel(nn.Module):
def sample(self, z_cond, uv_mask, noise=None):
if noise is None:
eps = torch.randn(z_cond.shape).to(uv_mask.device)
eps = torch.normal(mean=0.0, std=1.0, size=z_cond.shape).to(uv_mask.device)
else:
eps = noise
models/pyramid_model.py
+55 -2
View File
@@ -10,6 +10,29 @@ from models.pyramid_level import PyramidLevel
from torchvision.transforms import Resize
from torchvision.transforms.functional import InterpolationMode
# class SampleParallel(nn.Module):
# def __init__(self, levels):
# super().__init__()
# self.levels = nn.ModuleList(levels)
# def forward(self, z_cond, uv_masks, noise=None):
# level_outs = nn.parallel.data_parallel(self._forward_level, [(self.levels[i], z_cond, uv_masks[i], noise) for i in range(len(uv_masks))], dim=0, device_ids=[0, 0, 0, 0], output_device=0)
# # level_outs = nn.parallel.data_parallel(self.levels, [(z_cond, uv_masks[i], noise) for i in range(len(uv_masks))], device_ids=[0, 0, 0, 0], output_device=0)
# out = torch.zeros((uv_masks[0].shape[0], 3, 64, 32)).to(uv_masks[0].device)
# for i in range(len(self.levels)):
# level_out = level_outs[i]
# if i != len(self.levels) - 1:
# out = out + level_out
# out = self.levels[i].enhance(out, uv_masks[i])
# out = self.resizers[self.levels[i+1].size](out)
# out = out * uv_masks[i+1][:,None,:,:]
# else:
# out = out + level_out
# out = out * uv_masks[0][:,None,:,:]
# return out
# def _forward_level(self, level, z_cond, uv_mask, noise=None):
# return level(z_cond, uv_mask, noise=noise)
class Pyramid(nn.Module):
@@ -43,17 +66,24 @@ class Pyramid(nn.Module):
self.use_normals = config['use_normals']
def sample(self, uv_template, uv_body_posed, uv_normals, uv_mask, noise=None):
z_cond = self.encode_conds(uv_template, uv_body_posed, uv_normals)
# uv_masks = [uv_mask[64], uv_mask[128], uv_mask[256], uv_mask[512]]
# sampleP = SampleParallel(self.levels)
# out = sampleP(z_cond, uv_masks, noise=noise)
# return out
assert len(self.levels) == 4
out = torch.zeros((uv_template.shape[0], 3, 64, 32)).to(uv_template.device)
for i in range(len(self.levels)):
level_out = self.levels[i].sample(z_cond, uv_mask[self.levels[i].size], noise=noise)
level_out = self.levels[i].sample(z_cond, uv_mask[self.levels[i].size], noise=noise[i] if noise is not None else None)
if i != len(self.levels) - 1:
out = out + level_out
out = out + level_out
out = self.levels[i].enhance(out, uv_mask[self.levels[i].size])
out = self.resizers[self.levels[i+1].size](out)
out = out * uv_mask[self.levels[i+1].size][:,None,:,:]
@@ -63,7 +93,30 @@ class Pyramid(nn.Module):
out = out * uv_mask[512][:,None,:,:]
return out
def sample_level(self, uv_template, uv_body_posed, uv_normals, uv_mask, prev, level, noise=None):
z_cond = self.encode_conds(uv_template, uv_body_posed, uv_normals)
assert len(self.levels) == 4
i = level
out = prev
level_out = self.levels[i].sample(z_cond, uv_mask[self.levels[i].size], noise=noise[i] if noise is not None else None)
if i != len(self.levels) - 1:
out = out + level_out
out = self.levels[i].enhance(out, uv_mask[self.levels[i].size])
out = self.resizers[self.levels[i+1].size](out)
out = out * uv_mask[self.levels[i+1].size][:,None,:,:]
else:
out = out + level_out
out = out * uv_mask[512][:,None,:,:]
return out
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)