save

models/cvae_model.py

@@ -55,28 +55,24 @@ class CVAE(nn.Module):
         if self.avg_garment is not None:
             assert not self.use_offsets_over_body, 'When avg_garment is given, we must use skinning!'
 
-    # @tf.function
-    def sample(self, uv_template, uv_body_posed, uv_curv, uv_mask, noise=None, add_noise=False):
+    
+    def sample(self, uv_template, uv_body_posed, uv_normals, uv_mask, noise=None):
         if noise is None:
-            eps = tf.random.normal(shape=(uv_template.shape[0], 8, 4, self.latent_dim))
+            eps = torch.randn(uv_template.shape[0], self.latent_shape[0], self.latent_shape[1], self.latent_shape[2]).to(uv_template.device)
         else:
             eps = noise
 
-        z_cond = self.encode_conds(uv_template, uv_body_posed, uv_curv, training=False)
-        if add_noise:
-            z_cond = z_cond + (tf.random.uniform(z_cond.shape, maxval=self.noise_amount * 2) - self.noise_amount)
-        
-        z = tf.concat([eps, z_cond], axis=-1)
+        z_cond = self.encode_conds(uv_template, uv_body_posed, uv_normals)
+        z = torch.cat([eps, z_cond], dim=1)
+        x = self.decoder(z)
 
-        x = self.decoder(z, training=False)
-        
         if self.avg_garment is not None:
-            x = x + self.avg_garment[tf.newaxis,:,:,:]
-        
+            x = x + self.avg_garment[None,:,:,:]
+    
         if self.use_offsets_over_body:
             x = x + uv_body_posed
 
-        x = x * uv_mask[:,:,:,tf.newaxis]
+        x = x * uv_mask[:,None,:,:]
 
         return x
 
@@ -94,11 +90,11 @@ class CVAE(nn.Module):
         return out
 
 
-    def encode_conds(self, uv_template, uv_body_posed, uv_curv):
+    def encode_conds(self, uv_template, uv_body_posed, uv_normals):
         z_body = self.body_enc(uv_body_posed)
         z_static = self.static_enc(uv_template)
         if self.use_normals:
-            z_norm = self.norm_enc(uv_curv)
+            z_norm = self.norm_enc(uv_normals)
             z_cond = torch.concat([z_static, z_body, z_norm], dim=1)
         else:
             z_cond = torch.concat([z_static, z_body], dim=1)

models/pyramid_level.py

@@ -39,22 +39,16 @@ class PyramidLevel(nn.Module):
 
     def sample(self, z_cond, uv_mask, noise=None):
         if noise is None:
-            eps = torch.normal(shape=(uv_mask.shape[0], self.latent_dim))
+            eps = torch.randn(z_cond.shape).to(uv_mask.device)
         else:
             eps = noise
 
-        z = torch.concat([eps, z_cond], dim=-1)
-
+        z = torch.cat([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:,:]
+        x = x * uv_mask[:,None,:,:]
 
         return x
 
-
     def reparameterize(self, mean, logvar, fix_noise=False):
         if self.training == True:
             if not fix_noise:

models/pyramid_model.py

@@ -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)