gnn with weights

configs/config_abl_all_post.py

@@ -19,9 +19,9 @@ config['only_curvatures'] =           False
 ablation_count = [config['use_skinning'], config['remove_global_orientation'], config['use_curvature'], config['use_extra_losses'], config['use_offsets']].count(False)
 assert ablation_count == 1 or ablation_count == 0
 
-config['gpu'] = '4'
-config['name'] = 'iccv_weights_test_10'
-config['notes'] = 'smaller weight loss coeff, 1e-6, convolutional weight learning, add gnn'
+config['gpu'] = '5'
+config['name'] = 'iccv_weights_test_11'
+config['notes'] = 'smaller weight loss coeff, 1e-6, convolutional weight learning, add gnn but smaller and remove softmac from weights convnext'
 
 config['avg_garment_path'] = 'custom_data/train_no_shape_avg_uv.npy' #custom_data/train_no_shape_avg_uv.npy # train_no_shape_avg_uv  train_no_shape_one_hot_avg_uv
 if config['use_skinning']:
@@ -53,7 +53,7 @@ if config['is_cape']:
   config['smpl_res'] = 3638
 
 
-config['pretrained_model'] = 'checkpoints/post/iccv_weights_test_5'
+config['pretrained_model'] = ''
 
 ###############    
 config['lr'] = 0.0001#                 0.0001

convnext/models/convnext_tf.py

@@ -258,7 +258,8 @@ class ConvNeXtReverse(tf.keras.Model):
         if not learn_weights:
             self.last = Dense(3, activation='tanh')
         else:
-            self.last = Activation('softmax')
+            # self.last = Activation('softmax')
+            self.last = lambda x, training: x
             # self.last = Dense(24, activation='softmax')
 
 

cvae.py

@@ -281,7 +281,7 @@ def plot_everything(batch: dict, type: str):
     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()
+    mesh_weigths = smoother_weights(mesh_weigths, training=False)
     
     orig_img = uv_orig
     uv_static_img = batch['uv_static']
@@ -289,6 +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)]
+    # min max normalization of weights
+    weights = (weights - tf.reduce_min(weights, axis=[1,2,3], keepdims=True)) / (tf.reduce_max(weights, axis=[1,2,3], keepdims=True) - tf.reduce_min(weights, axis=[1,2,3], keepdims=True))
     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")
@@ -555,8 +557,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)
+    smoother = Smoother(features=24, activation_middle='tanh', activation_final='tanh', learn_weights=False)
+    smoother_weights = Smoother(features=3, activation_middle=None, activation_final='softmax', learn_weights=True)
 
     optimizer = tf.optimizers.Adam(learning_rate=config['lr'])
     balancing_activated_step = 0
@@ -581,6 +583,7 @@ if __name__ == '__main__':
 
         train_loss = train_step(batch)
 
+        if step % 25 == 0:
             data_valid = valid_iterator.get_next()
             data_valid = static_dataset_valid.expand_batch(batch_to_dict(data_valid, config['signature_dynamic']))
             valid_loss = valid_step(data_valid)

models/gnn.py

@@ -18,7 +18,7 @@ class Smoother(tf.keras.Model):
         self.a = sp_matrix_to_sp_tensor(sio.loadmat("custom_utils/ssmpl0.mat")['L'].tocoo())
         self.conv1 = GCNConv(features, activation=activation_middle, use_bias=False)
         self.conv2 = GCNConv(features, activation=activation_middle, use_bias=False)
-        self.conv3 = GCNConv(features, activation=activation_middle, use_bias=False)
+        # self.conv3 = GCNConv(features, activation=activation_middle, use_bias=False)
         # self.conv4 = GCNConv(64, activation='tanh', use_bias=False)
         # self.conv5 = GCNConv(64, activation='tanh', use_bias=False)
         self.conv6 = GCNConv(3 if not learn_weights else 24, activation=activation_final, use_bias=False)
@@ -27,7 +27,7 @@ class Smoother(tf.keras.Model):
         x = inputs
         x = self.conv1([x, self.a], training=training)
         x = self.conv2([x, self.a], training=training)
-        x = self.conv3([x, self.a], training=training)
+        # x = self.conv3([x, self.a], training=training)
         # x = self.conv4([x, self.a], training=training)
         # x = self.conv5([x, self.a], training=training)
         x = self.conv6([x, self.a], training=training)