weight test with gnn

cvae.py

@@ -42,6 +42,7 @@ from custom_utils.custom_utils import batch_to_dict, load_model, plot_3d, UV2Mes
 from dataset import DatasetFactory, StaticDataset
 from custom_utils.smpl import SMPL
 from models.cvae_model import CVAENew as CVAE
+from models.gnn import Smoother
 
 
 def init_wandb():
@@ -92,7 +93,9 @@ def train_step(batch):
         pred, mean, logvar, weights = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']], training=True)
 
         mesh_pred = uv2mesh(pred)
+        mesh_pred = smoother(mesh_pred, training=True)
         mesh_weigths = uv2mesh_label(weights)
+        mesh_weigths = smoother_weights(mesh_weigths, training=True)
 
         if config['use_skinning']:
             shape = batch['pose_shape'][:, 72:]
@@ -137,8 +140,8 @@ def train_step(batch):
 
     # ms_ssim = tf.reduce_mean(tf.image.ssim_multiscale(pred + 5, batch['uv_unposed'] + 5, max_val=10, power_factors=(0.5, 0.2, 0.1, 0.1, 0.1)))
     metric = tf.reduce_sum(tf.math.sqrt(tf.reduce_sum(batch['mesh_mask'][:,:,tf.newaxis] * (mesh_orig - mesh_pred_posed) ** 2, axis=-1))) / tf.reduce_sum(batch['mesh_mask'])
-    gradients = tape.gradient(loss, model.trainable_variables)
-    optimizer.apply_gradients(zip(gradients, model.trainable_variables))
+    gradients = tape.gradient(loss, model.trainable_variables + smoother.trainable_variables + smoother_weights.trainable_variables)
+    optimizer.apply_gradients(zip(gradients, model.trainable_variables + smoother.trainable_variables + smoother_weights.trainable_variables))
     return { 'train/loss': loss, 'train/edge_loss': e_loss, 'train/collision_loss': c_loss, 'train/euclid_error': metric,
              'train/reconstr_loss': reconstr_loss, 'train/collision_ratio': c_ratio,
              'train/mean': mean, 'train/logvar': logvar, 'train/kl_loss': kl, 'train/var': np.e ** logvar,
@@ -157,7 +160,9 @@ def valid_step(batch):
     pred, mean, logvar, weights = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']], training=False)
 
     mesh_pred = uv2mesh(pred)
+    mesh_pred = smoother(mesh_pred, training=False)
     mesh_weigths = uv2mesh_label(weights)
+    mesh_weigths = smoother_weights(mesh_weigths, training=False)
 
     if config['use_skinning']:
         shape = batch['pose_shape'][:, 72:]
@@ -273,8 +278,10 @@ def plot_everything(batch: dict, type: str):
 
     pred, mean, logvar, weights = model([uv_orig, batch['uv_static'], batch['uv_body_posed'], batch['uv_curvature'], batch['uv_mask']], training=False)
     
-    mesh_pred = uv2mesh(pred).numpy()
+    mesh_pred = uv2mesh(pred)
+    mesh_pred = smoother(mesh_pred, training=False).numpy()
     mesh_weigths = uv2mesh_label(weights)
+    mesh_weigths = smoother_weights(mesh_weigths, training=False).numpy()
     
     orig_img = uv_orig
     uv_static_img = batch['uv_static']
@@ -282,7 +289,8 @@ def plot_everything(batch: dict, type: str):
     pred_img = pred
     img = tf.concat((uv_static_img, uv_body_posed_img, orig_img, pred_img), axis=2)
     imgs = [wandb.Image((img[i] + 1) / 2.0) for i in range(4)]
-    
+    weights_imgs = [wandb.Image(np.concatenate((np.concatenate([weights[b,:,:,i].numpy() for i in range(12)], axis=-1), np.concatenate([weights[b,:,:,i].numpy() for i in range(12,24)], axis=-1)), axis=0)) for b in range(4)]
+
     F = np.load("custom_utils/smpl_faces" + skirt_modifier + ("_orig" if config['is_cape'] else "") + ".npy")
     mask = batch['mesh_mask'].numpy().astype(bool)
 
@@ -332,7 +340,7 @@ def plot_everything(batch: dict, type: str):
     mean_hist = wandb.Histogram(sequence=mean.numpy().mean(axis=0))
     logvar_hist = wandb.Histogram(sequence=logvar.numpy().mean(axis=0))
 
-    return { **figs, **figs_coll, **figs_weights, **{ type + '/images': imgs, type + '/mean_hist': mean_hist, type + '/logvar_hist': logvar_hist } }
+    return { **figs, **figs_coll, **figs_weights, **{ type + '/images': imgs, type + '/weight_images': weights_imgs, type + '/mean_hist': mean_hist, type + '/logvar_hist': logvar_hist } }
 
 
 def sample(folder): 
@@ -547,6 +555,8 @@ if __name__ == '__main__':
                                         is_cape=config['is_cape'])
 
     model = CVAE(config)
+    smoother = Smoother(features=16, activation_middle='tanh', activation_final='tanh', learn_weights=False)
+    smoother_weights = Smoother(features=16, activation_middle=None, activation_final='softmax', learn_weights=True)
 
     optimizer = tf.optimizers.Adam(learning_rate=config['lr'])
     balancing_activated_step = 0