SpatialTransfomer was embedded into a Keras model, and it is downloaded from the repo releases

DeepDeformationMapRegistration/layers/SpatialTransformer.py

@@ -3,6 +3,9 @@ import tensorflow.keras.backend as K
 import tensorflow as tf
 import neurite as ne
 
+import h5py
+from DeepDeformationMapRegistration.utils.constants import IMG_SHAPE, DISP_MAP_SHAPE
+
 
 class SpatialTransformer(kl.Layer):
     """
@@ -184,3 +187,15 @@ class SpatialTransformer(kl.Layer):
         # test single
         return ne.utils.interpn(vol, loc, interp_method=interp_method, fill_value=fill_value)
 
+
+if __name__ == "__main__":
+    output_file = './spatialtransformer.h5'
+
+    in_dm = tf.keras.Input(DISP_MAP_SHAPE)
+    in_image = tf.keras.Input(IMG_SHAPE)
+    pred = SpatialTransformer(interp_method='linear', indexing='ij', single_transform=False)([in_image, in_dm])
+
+    model = tf.keras.Model(inputs=[in_image, in_dm], outputs=pred)
+
+    model.save(output_file)
+    print(f"SpatialTransformer layer saved in: {output_file}")

DeepDeformationMapRegistration/main.py

@@ -19,7 +19,7 @@ import DeepDeformationMapRegistration.utils.constants as C
 from DeepDeformationMapRegistration.utils.nifti_utils import save_nifti
 from DeepDeformationMapRegistration.utils.operators import min_max_norm
 from DeepDeformationMapRegistration.utils.misc import resize_displacement_map
-from DeepDeformationMapRegistration.utils.model_utils import get_models_path, load_model
+from DeepDeformationMapRegistration.utils.model_utils import get_models_path, load_model, get_spatialtransformer_model
 from DeepDeformationMapRegistration.utils.logger import LOGGER
 
 from importlib.util import find_spec
@@ -279,8 +279,11 @@ def main():
 
     LOGGER.info(f'Getting 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]
+    ST_MODEL_FILE = get_spatialtransformer_model()
 
     network, registration_model = load_model(MODEL_FILE, False, True)
+    spatialtransformer_model = tf.keras.models.load_model(ST_MODEL_FILE,
+                                                          custom_objects={'SpatialTransformer': SpatialTransformer})
 
     LOGGER.info('Computing registration')
     with sess.as_default():
@@ -297,7 +300,8 @@ def main():
 
         LOGGER.info('Applying displacement map...')
         time_pred_img_start = time.time()
-        pred_image = SpatialTransformer(interp_method='linear', indexing='ij', single_transform=False)([moving_image[np.newaxis, ...], disp_map[np.newaxis, ...]]).eval()
+        # pred_image = SpatialTransformer(interp_method='linear', indexing='ij', single_transform=False)([moving_image[np.newaxis, ...], disp_map[np.newaxis, ...]]).eval()
+        pred_image = spatialtransformer_model.predict([moving_image[np.newaxis, ...], disp_map[np.newaxis, ...]])
         time_pred_img_end = time.time()
         LOGGER.info(f'\t... done ({time_pred_img_end - time_pred_img_start} s)')
         pred_image = pred_image[0, ...]

DeepDeformationMapRegistration/utils/model_utils.py

@@ -63,3 +63,16 @@ def load_model(weights_file_path: str, trainable: bool = False, return_registrat
         ret_val = (ret_val, ret_val.get_registration_model())
 
     return ret_val
+
+
+def get_spatialtransformer_model():
+    url = 'https://github.com/jpdefrutos/DDMR/releases/download/spatialtransformer_model_v0/spatialtransformer.h5'
+    file_path = os.path.join(os.getcwd(), 'models', 'spatialtransformer.h5')
+    if not os.path.exists(file_path):
+        LOGGER.info(f'Model not found. Downloading from {url}... ')
+        os.makedirs(os.path.split(file_path)[0], exist_ok=True)
+        download(url, file_path)
+        LOGGER.info(f'... downloaded model. Stored in {file_path}')
+    else:
+        LOGGER.info(f'Found model: {file_path}')
+    return file_path