Refactor the model utils functions (download weight files, get weight files path, and load models)

DeepDeformationMapRegistration/main.py

@@ -28,8 +28,7 @@ from DeepDeformationMapRegistration.losses import StructuralSimilarity_simplifie
 from DeepDeformationMapRegistration.ms_ssim_tf import MultiScaleStructuralSimilarity
 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 DeepDeformationMapRegistration.utils.model_utils import get_models_path, load_model
 
 from importlib.util import find_spec
 

DeepDeformationMapRegistration/networks.py

@@ -9,23 +9,6 @@ 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/{model_downloader.py → model_utils.py}

@@ -2,7 +2,8 @@ import os
 import requests
 from datetime import datetime
 from email.utils import parsedate_to_datetime, formatdate
-from DeepDeformationMapRegistration.utils.constants import ANATOMIES, MODEL_TYPES
+from DeepDeformationMapRegistration.utils.constants import ANATOMIES, MODEL_TYPES, ENCODER_FILTERS, DECODER_FILTERS, IMG_SHAPE
+import voxelmorph as vxm
 
 
 # taken from: https://lenon.dev/blog/downloading-and-caching-large-files-using-python/
@@ -35,9 +36,25 @@ def get_models_path(anatomy: str, model_type: str, output_root_dir: str):
     assert model_type in MODEL_TYPES.keys(), 'Invalid model type'
     anatomy = ANATOMIES[anatomy]
     model_type = MODEL_TYPES[model_type]
-    url = 'https://github.com/jpdefrutos/DDMR/releases/download/trained-models/' + anatomy + '_' + model_type + '.h5'
+    url = 'https://github.com/jpdefrutos/DDMR/releases/download/trained_models_v0/' + anatomy + '_' + model_type + '.h5'
     file_path = os.path.join(output_root_dir, 'models', anatomy, model_type + '.h5')
     if not os.path.exists(file_path):
         os.makedirs(os.path.split(file_path)[0], exist_ok=True)
         download(url, file_path)
     return file_path
+
+
+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