Merge pull request #45 from andreped/liver-support

Added seg-guiding support for demo + fixed CLI arg bug in main script

ddmr/main.py

@@ -171,9 +171,9 @@ def main():
     parser = argparse.ArgumentParser()
     parser.add_argument('-f', '--fixed', type=str, help='Path to fixed image file (NIfTI)')
     parser.add_argument('-m', '--moving', type=str, help='Path to moving segmentation image file (NIfTI)', default=None)
-    parser.add_argument('-F', '--fixedsegm', type=str, help='Path to fixed image segmentation file(NIfTI)',
+    parser.add_argument('-fs', '--fixedsegm', type=str, help='Path to fixed image segmentation file(NIfTI)',
                         default=None)
-    parser.add_argument('-M', '--movingsegm', type=str, help='Path to moving image file (NIfTI)')
+    parser.add_argument('-ms', '--movingsegm', type=str, help='Path to moving image file (NIfTI)')
     parser.add_argument('-o', '--outputdir', type=str, help='Output directory', default='./Registration_output')
     parser.add_argument('-a', '--anatomy', type=str, help='Anatomical structure: liver (L) (Default) or brain (B)',
                         default='L')

demo/src/compute.py

@@ -1,6 +1,11 @@
 import subprocess as sp
 
 
-def run_model(fixed_path, moving_path, output_path, task):
-    sp.check_call(["ddmr", "--fixed", fixed_path, "--moving", moving_path, \
-                   "-o", output_path, "-a", task, "--model", "BL-NS", "--original-resolution"])
+def run_model(fixed_path=None, moving_path=None, fixed_seg_path=None, moving_seg_path=None, output_path=None, task="B"):
+    if (fixed_seg_path is None) or (moving_seg_path is None):
+        print("The fixed or moving segmentation were not provided and are thus ignored for inference.")
+        sp.check_call(["ddmr", "-f", fixed_path, "-m", moving_path, \
+                    "-o", output_path, "-a", task, "--model", "BL-NS", "--original-resolution"])
+    else:
+        sp.check_call(["ddmr", "-f", fixed_path, "-m", moving_path, "-fs", fixed_seg_path, "-ms", moving_seg_path,
+                    "-o", output_path, "-a", task, "--model", "BL-NS", "--original-resolution"])

demo/src/gui.py

@@ -1,7 +1,4 @@
-import os
-
 import gradio as gr
-import numpy as np
 
 from .compute import run_model
 from .utils import load_ct_to_numpy
@@ -52,11 +49,22 @@ class WebUI:
         return [f.name for f in files]
 
     def process(self, mesh_file_names):
+        if not (len(mesh_file_names) in [2, 4]):
+            raise ValueError("Unsupported number of elements were provided as input to the DDMR CLI."
+                             "Either provided 2 or 4 elements, where the two first being the fixed"
+                             "and moving CT/MRIs and the two other being the binary segmentation"
+                             "which will be used for ROI filtering in preprocessing.")
         fixed_image_path = mesh_file_names[0].name
         moving_image_path = mesh_file_names[1].name
         output_path = self.cwd
 
-        run_model(fixed_image_path, moving_image_path, output_path, self.class_names[self.class_name])
+        if len(mesh_file_names) == 2:
+            run_model(fixed_image_path, moving_image_path, output_path, self.class_names[self.class_name])
+        else:
+            fixed_seg_path = mesh_file_names[2].name
+            moving_seg_path = mesh_file_names[3].name
+            
+            run_model(fixed_image_path, moving_image_path, fixed_seg_path, moving_seg_path, output_path, self.class_names[self.class_name])
 
         self.fixed_images = load_ct_to_numpy(fixed_image_path)
         self.moving_images = load_ct_to_numpy(moving_image_path)