Created convenient function for loading models

DeepDeformationMapRegistration/main.py

@@ -7,9 +7,9 @@ import subprocess
 import logging
 import time
 
-currentdir = os.path.dirname(os.path.realpath(__file__))
-parentdir = os.path.dirname(currentdir)
-sys.path.append(parentdir)  # PYTHON > 3.3 does not allow relative referencing
+# currentdir = os.path.dirname(os.path.realpath(__file__))
+# parentdir = os.path.dirname(currentdir)
+# sys.path.append(parentdir)  # PYTHON > 3.3 does not allow relative referencing
 
 import tensorflow as tf
 
@@ -29,6 +29,7 @@ from DeepDeformationMapRegistration.ms_ssim_tf import MultiScaleStructuralSimila
 from DeepDeformationMapRegistration.utils.operators import min_max_norm
 from DeepDeformationMapRegistration.utils.misc import resize_displacement_map
 from DeepDeformationMapRegistration.utils.model_downloader import get_models_path
+from DeepDeformationMapRegistration.networks import load_model
 
 from importlib.util import find_spec
 
@@ -284,39 +285,7 @@ def main():
     LOGGER.info(f'Using model: {"Brain" if args.anatomy == "B" else "Liver"} -> {args.model}')
     MODEL_FILE = get_models_path(args.anatomy, args.model, os.getcwd())  # MODELS_FILE[args.anatomy][args.model]
 
-    # try:
-    #     network = tf.keras.models.load_model(MODEL_FILE,
-    #                                          {'VxmDense': vxm.networks.VxmDense,
-    #                                           # 'VxmDenseSemiSupervisedSeg': vxm.networks.VxmDenseSemiSupervisedSeg,
-    #                                           'AdamAccumulated': AdamAccumulated
-    #                                           },
-    #                                          compile=False)
-    # except ValueError as e:
-    #     enc_features = [32, 64, 128, 256, 512, 1024]  # const.ENCODER_FILTERS
-    #     dec_features = enc_features[::-1] + [16, 16]  # const.ENCODER_FILTERS[::-1]
-    #     nb_features = [enc_features, dec_features]
-    #     if re.search('^UW|SEGGUIDED_', MODEL_FILE):
-    #         network = vxm.networks.VxmDense(inshape=IMAGE_INTPUT_SHAPE[:-1],
-    #                                                    nb_unet_features=nb_features,
-    #                                                    int_steps=0,
-    #                                                    int_downsize=1,
-    #                                                    seg_downsize=1)
-    #     else:
-    #         network = vxm.networks.VxmDense(inshape=IMAGE_INTPUT_SHAPE[:-1],
-    #                                                    nb_unet_features=nb_features,
-    #                                                    int_steps=0)
-    #     network.load_weights(MODEL_FILE, by_name=True)
-
-    enc_features = [32, 64, 128, 256, 512, 1024]  # const.ENCODER_FILTERS
-    dec_features = enc_features[::-1] + [16, 16]  # const.ENCODER_FILTERS[::-1]
-    nb_features = [enc_features, dec_features]
-    network = vxm.networks.VxmDense(inshape=C.IMG_SHAPE[:-1],
-                                    nb_unet_features=nb_features,
-                                    int_steps=0)
-    network.load_weights(MODEL_FILE, by_name=True)
-    network.trainable = False
-
-    registration_model = network.get_registration_model()
+    network, registration_model = load_model(MODEL_FILE, False, True)
     deb_model = network.apply_transform
 
     LOGGER.info('Computing registration')

DeepDeformationMapRegistration/networks.py

@@ -1,14 +1,31 @@
 import os, sys
-currentdir = os.path.dirname(os.path.realpath(__file__))
-parentdir = os.path.dirname(currentdir)
-sys.path.append(parentdir)  # PYTHON > 3.3 does not allow relative referencing
-
-PYCHARM_EXEC = os.getenv('PYCHARM_EXEC') == 'True'
+# currentdir = os.path.dirname(os.path.realpath(__file__))
+# parentdir = os.path.dirname(currentdir)
+# sys.path.append(parentdir)  # PYTHON > 3.3 does not allow relative referencing
+#
+# PYCHARM_EXEC = os.getenv('PYCHARM_EXEC') == 'True'
 
 import tensorflow as tf
 import voxelmorph as vxm
 from voxelmorph.tf.modelio import LoadableModel, store_config_args
 from tensorflow.keras.layers import UpSampling3D
+from DeepDeformationMapRegistration.utils.constants import ENCODER_FILTERS, DECODER_FILTERS, IMG_SHAPE
+
+
+def load_model(weights_file_path: str, trainable: bool = False, return_registration_model: bool=True):
+    assert os.path.exists(weights_file_path), f'File {weights_file_path} not found'
+    assert weights_file_path.endswith('h5'), 'Invalid file extension. Expected .h5'
+
+    ret_val = vxm.networks.VxmDense(inshape=IMG_SHAPE[:-1],
+                                    nb_unet_features=[ENCODER_FILTERS, DECODER_FILTERS],
+                                    int_steps=0)
+    ret_val.load_weights(weights_file_path, by_name=True)
+    ret_val.trainable = trainable
+
+    if return_registration_model:
+        ret_val = (ret_val, ret_val.get_registration_model())
+
+    return ret_val
 
 
 class WeaklySupervised(LoadableModel):

DeepDeformationMapRegistration/utils/constants.py

@@ -196,8 +196,8 @@ DROPOUT = True
 DROPOUT_RATE = 0.2
 MAX_DATA_SIZE = (1000, 1000, 1)
 PLATEAU_THR = 0.01  # A slope between +-PLATEAU_THR will be considered a plateau for the LR updating function
-ENCODER_FILTERS = [4, 8, 16, 32, 64]
-
+ENCODER_FILTERS = [32, 64, 128, 256, 512, 1024]
+DECODER_FILTERS = ENCODER_FILTERS[::-1] + [16, 16]
 # SSIM
 SSIM_FILTER_SIZE = 11  # Size of Gaussian filter
 SSIM_FILTER_SIGMA = 1.5  # Width of Gaussian filter
@@ -205,7 +205,7 @@ SSIM_K1 = 0.01  # Def. 0.01
 SSIM_K2 = 0.03  # Recommended values 0 < K2 < 0.4
 MAX_VALUE = 1.0  # Maximum intensity values
 
-# Mathematic constants
+# Mathematics constants
 EPS = 1e-8
 EPS_tf = tf.constant(EPS, dtype=tf.float32)
 LOG2 = tf.math.log(tf.constant(2, dtype=tf.float32))