Major refactor on deployment

.gitignore

@@ -2,3 +2,8 @@
 .vs/
 *__pycache__/
 venv/
+*.nii
+*.nii.gz
+*.egg-info/
+*.h5
+*.log

DeepDeformationMapRegistration/main.py

@@ -1,16 +1,10 @@
 import datetime
 import os, sys
 import shutil
-import re
 import argparse
 import subprocess
 import logging
 import time
-import warnings
-
-# 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
 
@@ -20,13 +14,10 @@ from scipy.ndimage import gaussian_filter, zoom
 from skimage.measure import regionprops
 import SimpleITK as sitk
 
-import voxelmorph as vxm
 from voxelmorph.tf.layers import SpatialTransformer
 
 import DeepDeformationMapRegistration.utils.constants as C
 from DeepDeformationMapRegistration.utils.nifti_utils import save_nifti
-from DeepDeformationMapRegistration.losses import StructuralSimilarity_simplified, NCC
-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_utils import get_models_path, load_model
@@ -191,7 +182,7 @@ def main():
     parser.add_argument('--gpu', type=int,
                         help='In case of multi-GPU systems, limits the execution to the defined GPU number',
                         default=None)
-    parser.add_argument('--model', type=str, help='Which model to use: BL-N, BL-S, SG-ND, SG-NSD, UW-NSD, UW-NSDH',
+    parser.add_argument('--model', type=str, help='Which model to use: BL-N, BL-S, BL-NS, SG-ND, SG-NSD, UW-NSD, UW-NSDH',
                         default='UW-NSD')
     parser.add_argument('-d', '--debug', action='store_true', help='Produce additional debug information', default=False)
     parser.add_argument('-c', '--clear-outputdir', action='store_true', help='Clear output folder if this has content', default=False)
@@ -206,17 +197,7 @@ def main():
     assert os.path.exists(args.moving), 'Moving image not found'
     assert args.model in C.MODEL_TYPES.keys(), 'Invalid model type'
     assert args.anatomy in C.ANATOMIES.keys(), 'Invalid anatomy option'
-    if os.path.exists(args.outputdir) and len(os.listdir(args.outputdir)):
-        if args.clear_outputdir:
-            erase = 'y'
-        else:
-            erase = input('Output directory is not empty, erase content? (y/n)')
-        if erase.lower() in ['y', 'yes']:
-            shutil.rmtree(args.outputdir, ignore_errors=True)
-            print('Erased directory: ' + args.outputdir)
-        elif erase.lower() in ['n', 'no']:
-            args.outputdir = os.path.join(args.outputdir, datetime.datetime.now().strftime('%H%M%S_%Y%m%d'))
-            print('New output directory: ' + args.outputdir)
+
     os.makedirs(args.outputdir, exist_ok=True)
 
     log_format = '%(asctime)s [%(levelname)s]:\t%(message)s'
@@ -296,20 +277,11 @@ def main():
 
     # 3. Build the whole graph
     LOGGER.info('Building TF graph')
-    ### METRICS GRAPH ###
-    fix_img_ph = tf.compat.v1.placeholder(tf.float32, (1, None, None, None, 1), name='fix_img')
-    pred_img_ph = tf.compat.v1.placeholder(tf.float32, (1, None, None, None, 1), name='pred_img')
-
-    ssim_tf = StructuralSimilarity_simplified(patch_size=2, dim=3, dynamic_range=1.).metric(fix_img_ph, pred_img_ph)
-    ncc_tf = NCC(image_shape_or).metric(fix_img_ph, pred_img_ph)
-    mse_tf = vxm.losses.MSE().loss(fix_img_ph, pred_img_ph)
-    ms_ssim_tf = MultiScaleStructuralSimilarity(max_val=1., filter_size=3).metric(fix_img_ph, pred_img_ph)
 
     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]
 
     network, registration_model = load_model(MODEL_FILE, False, True)
-    deb_model = network.apply_transform
 
     LOGGER.info('Computing registration')
     with sess.as_default():
@@ -317,20 +289,17 @@ def main():
             registration_model.summary(line_length=C.SUMMARY_LINE_LENGTH)
         LOGGER.info('Computing displacement map...')
         time_disp_map_start = time.time()
-        # disp_map = registration_model.predict([moving_image[np.newaxis, ...], fixed_image[np.newaxis, ...]])
         p, disp_map = network.predict([moving_image[np.newaxis, ...], fixed_image[np.newaxis, ...]])
         time_disp_map_end = time.time()
         LOGGER.info(f'\t... done ({time_disp_map_end - time_disp_map_start})')
         disp_map = np.squeeze(disp_map)
         debug_save_image(np.squeeze(disp_map), 'disp_map_0_raw', args.outputdir, args.debug)
         debug_save_image(p[0, ...], 'img_4_net_pred_image', args.outputdir, args.debug)
-        # pred_image = min_max_norm(pred_image)
-        # pred_image_isot = zoom(pred_image[0, ...], zoom_factors, order=3)[np.newaxis, ...]
-        # fixed_image_isot = zoom(fixed_image[0, ...], zoom_factors, order=3)[np.newaxis, ...]
 
         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 = np.zeros_like(moving_image[np.newaxis, ...])  # @TODO: Replace this with Keras' Model with SpatialTransformer Layer
         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, ...]
@@ -340,31 +309,10 @@ def main():
             moving_image = moving_image_or
             fixed_image = fixed_image_or
             # disp_map = disp_map_or
-            pred_image = zoom(pred_image, 1/zoom_factors)
+            pred_image = zoom(pred_image, 1 / zoom_factors)
             pred_image = pad_crop_to_original_shape(pred_image, fixed_image_or.shape, crop_min)
             LOGGER.info('Done...')
 
-        LOGGER.info('Computing metrics...')
-        if args.original_resolution:
-            ssim, ncc, mse, ms_ssim = sess.run([ssim_tf, ncc_tf, mse_tf, ms_ssim_tf],
-                                               {'fix_img:0': fixed_image[np.newaxis,
-                                                                         crop_min[0]: crop_max[0],
-                                                                         crop_min[1]: crop_max[1],
-                                                                         crop_min[2]: crop_max[2],
-                                                                         ...],
-                                                'pred_img:0': pred_image[np.newaxis,
-                                                                         crop_min[0]: crop_max[0],
-                                                                         crop_min[1]: crop_max[1],
-                                                                         crop_min[2]: crop_max[2],
-                                                                         ...]})  # to only compare the deformed region!
-        else:
-
-            ssim, ncc, mse, ms_ssim = sess.run([ssim_tf, ncc_tf, mse_tf, ms_ssim_tf],
-                                               {'fix_img:0': fixed_image[np.newaxis, ...],
-                                                'pred_img:0': pred_image[np.newaxis, ...]})
-        ssim = np.mean(ssim)
-        ms_ssim = ms_ssim[0]
-
         if args.original_resolution:
             save_nifti(pred_image, os.path.join(args.outputdir, 'pred_image.nii.gz'), header=moving_image_header)
         else:
@@ -389,26 +337,11 @@ def main():
                 np.savez_compressed(os.path.join(args.outputdir, 'displacement_map.npz'), disp_map)
             else:
                 np.savez_compressed(os.path.join(os.path.join(args.outputdir, 'debug'), 'displacement_map.npz'), disp_map)
+        
         LOGGER.info('Predicted image and displacement map saved in: '.format(args.outputdir))
         LOGGER.info(f'Displacement map prediction time: {time_disp_map_end - time_disp_map_start} s')
         LOGGER.info(f'Predicted image time: {time_pred_img_end - time_pred_img_start} s')
 
-        LOGGER.info('Similarity metrics\n------------------')
-        LOGGER.info('SSIM: {:.03f}'.format(ssim))
-        LOGGER.info('NCC: {:.03f}'.format(ncc))
-        LOGGER.info('MSE: {:.03f}'.format(mse))
-        LOGGER.info('MS SSIM: {:.03f}'.format(ms_ssim))
-
-        # ssim, ncc, mse, ms_ssim = sess.run([ssim_tf, ncc_tf, mse_tf, ms_ssim_tf],
-        #                                    {'fix_img:0': fixed_image[np.newaxis, ...], 'pred_img:0': p})
-        # ssim = np.mean(ssim)
-        # ms_ssim = ms_ssim[0]
-        # LOGGER.info('\nSimilarity metrics (ROI)\n------------------')
-        # LOGGER.info('SSIM: {:.03f}'.format(ssim))
-        # LOGGER.info('NCC: {:.03f}'.format(ncc))
-        # LOGGER.info('MSE: {:.03f}'.format(mse))
-        # LOGGER.info('MS SSIM: {:.03f}'.format(ms_ssim))
-
     del registration_model
     LOGGER.info('Done')
     exit(0)

DeepDeformationMapRegistration/utils/logger.py

@@ -0,0 +1,4 @@
+import logging
+
+
+LOGGER = logging.getLogger(__name__)

Dockerfile

@@ -0,0 +1,63 @@
+# read the doc: https://huggingface.co/docs/hub/spaces-sdks-docker
+# you will also find guides on how best to write your Dockerfile
+FROM python:3.8-slim
+
+# set language, format and stuff
+ENV LANG=C.UTF-8 LC_ALL=C.UTF-8
+
+WORKDIR /code
+
+RUN apt-get update -y
+#RUN apt-get install -y python3 python3-pip
+RUN apt install git --fix-missing -y
+RUN apt install wget -y
+
+# installing other libraries
+RUN apt-get install python3-pip -y && \
+    apt-get -y install sudo
+RUN apt-get install curl -y
+RUN apt-get install nano -y
+RUN apt-get update && apt-get install -y git
+RUN apt-get install libblas-dev -y && apt-get install liblapack-dev -y
+RUN apt-get install gfortran -y
+RUN apt-get install libpng-dev -y
+RUN apt-get install python3-dev -y
+
+WORKDIR /code
+
+# install dependencies
+COPY ./demo/requirements.txt /code/demo/requirements.txt
+RUN pip install --no-cache-dir --upgrade -r /code/demo/requirements.txt
+
+# resolve issue with tf==2.4 and gradio dependency collision issue
+RUN pip install --force-reinstall typing_extensions==4.7.1
+
+# lower pydantic version to work with typing_extensions deprecation
+RUN pip install --force-reinstall "pydantic<2.0.0"
+
+# Install wget
+RUN apt install wget -y && \
+    apt install unzip
+
+# Set up a new user named "user" with user ID 1000
+RUN useradd -m -u 1000 user
+
+# Switch to the "user" user
+USER user
+
+# Set home to the user's home directory
+ENV HOME=/home/user \
+    PATH=/home/user/.local/bin:$PATH
+
+# Set the working directory to the user's home directory
+WORKDIR $HOME/app
+
+# Copy the current directory contents into the container at $HOME/app setting the owner to the user
+COPY --chown=user . $HOME/app
+
+# download test data
+RUN wget "https://github.com/jpdefrutos/DDMR/releases/download/test_data_brain_v0/ixi_image.nii.gz" && \
+    wget "https://github.com/jpdefrutos/DDMR/releases/download/test_data_brain_v0/ixi_image2.nii.gz"
+
+# CMD ["/bin/bash"]
+CMD ["python3", "app.py"]

README.md

@@ -1,3 +1,15 @@
+---
+title: 'DDMR: Deep Deformation Map Registration of CT/MRIs'
+colorFrom: indigo
+colorTo: indigo
+sdk: docker
+app_port: 7860
+emoji: 🧠
+pinned: false
+license: mit
+app_file: demo/app.py
+---
+
 <div align="center">
     <img src="https://user-images.githubusercontent.com/30429725/204778476-4d24c659-9287-48b8-b616-92016ffcf4f6.svg" alt="drawing" width="600">
 </div>
@@ -6,6 +18,8 @@
 
 <h1 align="center">DDMR: Deep Deformation Map Registration</h1>
 <h3 align="center">Learning deep abdominal CT registration through adaptive loss weighting and synthetic data generation</h3>
+ 
+# ⚠️***WARNING: Under construction*** 
 
 **DDMR** was developed by SINTEF Health Research. The corresponding manuscript describing the framework has been published in [PLOS ONE](https://journals.plos.org/plosone/) and the manuscript is openly available [here](https://journals.plos.org/plosone/article?id=10.1371/journal.pone.0282110).
 
@@ -22,8 +36,27 @@ source venv/bin/activate
 
 2. Install requirements:
 ```
-pip install -r requirements.txt
+pip install /path/to/clone/.
+```
+
+## 🤖 How to use
+Use the following CLI command to register images
+```
+ddmr --fixed path/to/fixed_image.nii.gz --moving path/to/moving_image.nii.gz --outputdir path/to/output/dir -a <anatomy> --model <model> --gpu <gpu-number> --original-resolution
 ```
+where:
+* anatomy: is the type of anatomy you want to register: B (brain) or L (liver)
+* model: is the model you want to use:
+    + BL-N (baseline with NCC)
+    + BL-NS (baseline with NCC and SSIM)
+    + SG-ND (segmentation guided with NCC and DSC)
+    + SG-NSD (segmentation guided with NCC, SSIM, and DSC)
+    + UW-NSD (uncertainty weighted with NCC, SSIM, and DSC)
+    + UW-NSDH (uncertainty weighted with NCC, SSIM, DSC, and HD).
+* gpu: is the GPU number you want to the model to run on, if you have multiple and want to use only one GPU
+* original-resolution: (flag) whether to upsample the registered image to the fixed image resolution (disabled if the flag is not present)
+
+Use ```ddmr --help``` to see additional options like using precomputed segmentations to crop the images to the desired ROI, or debugging.
 
 ## 🏋️‍♂️ Training
 

demo/app.py

@@ -0,0 +1,40 @@
+import os
+from argparse import ArgumentParser
+
+
+def main():
+    parser = ArgumentParser()
+    parser.add_argument(
+        "--cwd",
+        type=str,
+        default="/home/user/app/",
+        help="Set current working directory (path to app.py).",
+    )
+    parser.add_argument(
+        "--share",
+        type=int,
+        default=1,
+        help="Whether to enable the app to be accessible online"
+        "-> setups a public link which requires internet access.",
+    )
+    args = parser.parse_args()
+
+    print("Current working directory:", args.cwd)
+
+    if not os.path.exists(args.cwd):
+        raise ValueError("Chosen 'cwd' is not a valid path!")
+    if args.share not in [0, 1]:
+        raise ValueError(
+            "The 'share' argument can only be set to 0 or 1, but was:",
+            args.share,
+        )
+
+    # initialize and run app
+    print("Launching demo...")
+    from src.gui import WebUI
+    app = WebUI(cwd=args.cwd, share=args.share)
+    app.run()
+
+
+if __name__ == "__main__":
+    main()

demo/requirements.txt

@@ -0,0 +1,2 @@
+ddmr @ git+https://github.com/jpdefrutos/DDMR.git
+gradio==3.44.4

demo/src/compute.py

@@ -0,0 +1,3 @@
+def run_model(input_path):
+    from lungtumormask import mask
+    mask.mask(input_path, "./prediction.nii.gz", lung_filter=True, threshold=0.5, radius=1, batch_size=1)

demo/src/gui.py

@@ -0,0 +1,177 @@
+import os
+
+import gradio as gr
+
+from .inference import run_model
+from .utils import load_ct_to_numpy
+from .utils import load_pred_volume_to_numpy
+from .utils import nifti_to_glb
+
+
+class WebUI:
+    def __init__(
+        self,
+        model_name: str = None,
+        cwd: str = "/home/user/app/",
+        share: int = 1,
+    ):
+        # global states
+        self.images = []
+        self.pred_images = []
+
+        # @TODO: This should be dynamically set based on chosen volume size
+        self.nb_slider_items = 150
+
+        self.model_name = model_name
+        self.cwd = cwd
+        self.share = share
+
+        self.class_name = "meningioma"  # default
+        self.class_names = {
+            "meningioma": "MRI_Meningioma",
+            "low-grade": "MRI_LGGlioma",
+            "metastasis": "MRI_Metastasis",
+            "high-grade": "MRI_GBM",
+            "brain": "MRI_Brain",
+        }
+
+        self.result_names = {
+            "meningioma": "Tumor",
+            "low-grade": "Tumor",
+            "metastasis": "Tumor",
+            "high-grade": "Tumor",
+            "brain": "Brain",
+        }
+
+        # define widgets not to be rendered immediantly, but later on
+        self.slider = gr.Slider(
+            1,
+            self.nb_slider_items,
+            value=1,
+            step=1,
+            label="Which 2D slice to show",
+        )
+        self.volume_renderer = gr.Model3D(
+            clear_color=[0.0, 0.0, 0.0, 0.0],
+            label="3D Model",
+            visible=True,
+            elem_id="model-3d",
+        ).style(height=512)
+
+    def set_class_name(self, value):
+        print("Changed task to:", value)
+        self.class_name = value
+
+    def combine_ct_and_seg(self, img, pred):
+        return (img, [(pred, self.class_name)])
+
+    def upload_file(self, file):
+        return file.name
+
+    def process(self, mesh_file_name):
+        path = mesh_file_name.name
+        run_model(
+            path,
+            model_path=os.path.join(self.cwd, "resources/models/"),
+            task=self.class_names[self.class_name],
+            name=self.result_names[self.class_name],
+        )
+        nifti_to_glb("prediction.nii.gz")
+
+        self.images = load_ct_to_numpy(path)
+        self.pred_images = load_pred_volume_to_numpy("./prediction.nii.gz")
+        return "./prediction.obj"
+
+    def get_img_pred_pair(self, k):
+        k = int(k) - 1
+        out = [gr.AnnotatedImage.update(visible=False)] * self.nb_slider_items
+        out[k] = gr.AnnotatedImage.update(
+            self.combine_ct_and_seg(self.images[k], self.pred_images[k]),
+            visible=True,
+        )
+        return out
+
+    def run(self):
+        css = """
+        #model-3d {
+        height: 512px;
+        }
+        #model-2d {
+        height: 512px;
+        margin: auto;
+        }
+        #upload {
+        height: 120px;
+        }
+        """
+        with gr.Blocks(css=css) as demo:
+            with gr.Row():
+                file_output = gr.File(file_count="single", elem_id="upload")
+                file_output.upload(self.upload_file, file_output, file_output)
+
+                model_selector = gr.Dropdown(
+                    list(self.class_names.keys()),
+                    label="Task",
+                    info="Which task to perform - one model for"
+                    "each brain tumor type and brain extraction",
+                    multiselect=False,
+                    size="sm",
+                )
+                model_selector.input(
+                    fn=lambda x: self.set_class_name(x),
+                    inputs=model_selector,
+                    outputs=None,
+                )
+
+                run_btn = gr.Button("Run analysis").style(
+                    full_width=False, size="lg"
+                )
+                run_btn.click(
+                    fn=lambda x: self.process(x),
+                    inputs=file_output,
+                    outputs=self.volume_renderer,
+                )
+
+            with gr.Row():
+                gr.Examples(
+                    examples=[
+                        os.path.join(self.cwd, "RegLib_C01_1.nii"),
+                        os.path.join(self.cwd, "RegLib_C01_2.nii"),
+                    ],
+                    inputs=file_output,
+                    outputs=file_output,
+                    fn=self.upload_file,
+                    cache_examples=True,
+                )
+
+            with gr.Row():
+                with gr.Box():
+                    with gr.Column():
+                        image_boxes = []
+                        for i in range(self.nb_slider_items):
+                            visibility = True if i == 1 else False
+                            t = gr.AnnotatedImage(
+                                visible=visibility, elem_id="model-2d"
+                            ).style(
+                                color_map={self.class_name: "#ffae00"},
+                                height=512,
+                                width=512,
+                            )
+                            image_boxes.append(t)
+
+                        self.slider.input(
+                            self.get_img_pred_pair, self.slider, image_boxes
+                        )
+
+                        self.slider.render()
+
+                with gr.Box():
+                    self.volume_renderer.render()
+
+        # sharing app publicly -> share=True:
+        # https://gradio.app/sharing-your-app/
+        # inference times > 60 seconds -> need queue():
+        # https://github.com/tloen/alpaca-lora/issues/60#issuecomment-1510006062
+        demo.queue().launch(
+            server_name="0.0.0.0", server_port=7860, share=self.share
+        )

demo/src/utils.py

@@ -0,0 +1,67 @@
+import nibabel as nib
+import numpy as np
+from nibabel.processing import resample_to_output
+from skimage.measure import marching_cubes
+
+
+def load_ct_to_numpy(data_path):
+    if type(data_path) != str:
+        data_path = data_path.name
+
+    image = nib.load(data_path)
+    resampled = resample_to_output(image, None, order=0)
+    data = resampled.get_fdata()
+
+    data = np.rot90(data, k=1, axes=(0, 1))
+
+    data[data < -150] = -150
+    data[data > 250] = 250
+
+    data = data - np.amin(data)
+    data = data / np.amax(data) * 255
+    data = data.astype("uint8")
+
+    print(data.shape)
+    return [data[..., i] for i in range(data.shape[-1])]
+
+
+def load_pred_volume_to_numpy(data_path):
+    if type(data_path) != str:
+        data_path = data_path.name
+
+    image = nib.load(data_path)
+    resampled = resample_to_output(image, None, order=0)
+    data = resampled.get_fdata()
+
+    data = np.rot90(data, k=1, axes=(0, 1))
+
+    data[data > 0] = 1
+    data = data.astype("uint8")
+
+    print(data.shape)
+    return [data[..., i] for i in range(data.shape[-1])]
+
+
+def nifti_to_glb(path, output="prediction.obj"):
+    # load NIFTI into numpy array
+    image = nib.load(path)
+    resampled = resample_to_output(image, [1, 1, 1], order=1)
+    data = resampled.get_fdata().astype("uint8")
+
+    # extract surface
+    verts, faces, normals, values = marching_cubes(data, 0)
+    faces += 1
+
+    with open(output, "w") as thefile:
+        for item in verts:
+            thefile.write("v {0} {1} {2}\n".format(item[0], item[1], item[2]))
+
+        for item in normals:
+            thefile.write("vn {0} {1} {2}\n".format(item[0], item[1], item[2]))
+
+        for item in faces:
+            thefile.write(
+                "f {0}//{0} {1}//{1} {2}//{2}\n".format(
+                    item[0], item[1], item[2]
+                )
+            )

setup.py

@@ -23,21 +23,18 @@ setup(
     ],
     python_requires='>=3.6',
     install_requires=[
-        'fastrlock>=0.3',   # required by cupy-cuda110
-        'testresources',    # required by launchpadlib
         'scipy',
         'scikit-image',
         'simpleITK',
         'voxelmorph==0.1',
         'pystrum==0.1',
-        'tensorflow-gpu==1.14.0',
+        'tensorflow==2.13',
         'tensorflow-addons',
         'tensorflow-datasets',
         'tensorflow-metadata',
-        'tensorboard==1.14.0',
         'nibabel==3.2.1',
-        'numpy==1.18.5',
-        'h5py==2.10'
+        'numpy',
+        'h5py'
     ],
     entry_points={
         'console_scripts': ['ddmr=DeepDeformationMapRegistration.main:main']