save

2023-07-14 10:12:00 +02:00
parent bb16aaf868
commit 88531da784
16 changed files with 411 additions and 299 deletions
@@ -42,6 +42,28 @@ 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)
+
+        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)
+            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)