Updated call to DataGeneratorManager

TrainingScripts/Train_3d.py

@@ -15,6 +15,7 @@ from datetime import datetime
 
 import DeepDeformationMapRegistration.utils.constants as C
 from DeepDeformationMapRegistration.data_generator import DataGeneratorManager
+from DeepDeformationMapRegistration.losses import NCC
 from DeepDeformationMapRegistration.utils.misc import try_mkdir
 
 
@@ -44,19 +45,14 @@ vxm_model = vxm.networks.VxmDense(inshape=in_shape, nb_unet_features=nb_features
 
 
 # Losses and loss weights
-
-def comb_loss(y_true, y_pred):
-    return vxm.losses.MSE().loss(y_true, y_pred) + vxm.losses.NCC().loss(y_true, y_pred)
-
-
-losses = [comb_loss, vxm.losses.Grad('l2').loss]
+losses = [NCC(in_shape).loss, vxm.losses.Grad('l2').loss]
 loss_weights = [1., 0.01]
 
 # Compile the model
 vxm_model.compile(optimizer=tf.keras.optimizers.Adam(lr=1e-4), loss=losses, loss_weights=loss_weights)
 
 # Train
-output_folder = os.path.join('train_3d_mse_ncc_grad_'+datetime.now().strftime("%H%M%S-%d%m%Y"))
+output_folder = os.path.join('TrainingScripts/TrainOutput/baseline_LITS_NCC_'+datetime.now().strftime("%H%M%S-%d%m%Y"))
 try_mkdir(output_folder)
 try_mkdir(os.path.join(output_folder, 'checkpoints'))
 try_mkdir(os.path.join(output_folder, 'tensorboard'))

TrainingScripts/Train_3d_weaklySupervised.py

@@ -33,7 +33,9 @@ C.EPOCHS = 10000
 # Build data generator
 
 data_generator = DataGeneratorManager(C.TRAINING_DATASET, C.BATCH_SIZE, True, C.LIMIT_NUM_SAMPLES,
-                                      1 - C.TRAINING_PERC, voxelmorph=True, segmentations=True)
+                                      1 - C.TRAINING_PERC,
+                                      input_labels=[C.DG_LBL_MOV_VESSELS, C.DG_LBL_FIX_VESSELS, C.DG_LBL_MOV_IMG, C.DG_LBL_ZERO_GRADS],
+                                      output_labels=[])
 
 train_generator = data_generator.get_generator('train')
 validation_generator = data_generator.get_generator('validation')
@@ -72,7 +74,7 @@ full_model = tf.keras.Model(inputs=vxm_model.inputs + [grad], outputs=vxm_model.
 full_model.compile(optimizer=tf.keras.optimizers.Adam(lr=1e-4), loss=None)
 
 # Train
-output_folder = os.path.join('train_3d_multiloss_segm_haus_dice_ncc_grad_'+datetime.now().strftime("%H%M%S-%d%m%Y"))
+output_folder = os.path.join('TrainingScripts/TrainOutput/weaklysupervised_DCTHLN_UW_haus_dice_'+datetime.now().strftime("%H%M%S-%d%m%Y"))
 try_mkdir(output_folder)
 try_mkdir(os.path.join(output_folder, 'checkpoints'))
 try_mkdir(os.path.join(output_folder, 'tensorboard'))