Consolidate training and inference

FNNDSC · Apr 9, 2024 · dbbd8ac · dbbd8ac
1 parent 3c5f61f
commit dbbd8ac
Showing 1 changed file with 14 additions and 8 deletions.
diff --git a/spleenseg.py b/spleenseg.py
@@ -49,23 +49,29 @@
 
 DISPLAY_TITLE = r"""
 
-███████╗██████╗ ██╗     ███████╗███████╗███╗   ██╗███████╗███████╗ ██████╗      ████████╗██████╗  █████╗ ██╗███╗   ██╗
-██╔════╝██╔══██╗██║     ██╔════╝██╔════╝████╗  ██║██╔════╝██╔════╝██╔════╝      ╚══██╔══╝██╔══██╗██╔══██╗██║████╗  ██║
-███████╗██████╔╝██║     █████╗  █████╗  ██╔██╗ ██║███████╗█████╗  ██║  ███╗        ██║   ██████╔╝███████║██║██╔██╗ ██║
-╚════██║██╔═══╝ ██║     ██╔══╝  ██╔══╝  ██║╚██╗██║╚════██║██╔══╝  ██║   ██║        ██║   ██╔══██╗██╔══██║██║██║╚██╗██║
-███████║██║     ███████╗███████╗███████╗██║ ╚████║███████║███████╗╚██████╔╝███████╗██║   ██║  ██║██║  ██║██║██║ ╚████║
-╚══════╝╚═╝     ╚══════╝╚══════╝╚══════╝╚═╝  ╚═══╝╚══════╝╚══════╝ ╚═════╝ ╚══════╝╚═╝   ╚═╝  ╚═╝╚═╝  ╚═╝╚═╝╚═╝  ╚═══╝
+███████╗██████╗ ██╗     ███████╗███████╗███╗   ██╗███████╗███████╗ ██████╗
+██╔════╝██╔══██╗██║     ██╔════╝██╔════╝████╗  ██║██╔════╝██╔════╝██╔════╝
+███████╗██████╔╝██║     █████╗  █████╗  ██╔██╗ ██║███████╗█████╗  ██║  ███╗
+╚════██║██╔═══╝ ██║     ██╔══╝  ██╔══╝  ██║╚██╗██║╚════██║██╔══╝  ██║   ██║
+███████║██║     ███████╗███████╗███████╗██║ ╚████║███████║███████╗╚██████╔╝
+╚══════╝╚═╝     ╚══════╝╚══════╝╚══════╝╚═╝  ╚═══╝╚══════╝╚══════╝ ╚═════╝
 """
 
 
 parser = ArgumentParser(
     description="""
 A ChRIS DS plugin based on Project MONAI 3D Spleen Segmentation.
-This plugin performs the training component, with some gentle
+This plugin implements both training and inference, with some
 refactoring and pervasive type hinting.
     """,
     formatter_class=ArgumentDefaultsHelpFormatter,
 )
+parser.add_argument(
+    "--mode",
+    type=str,
+    default="training",
+    help="mode of behaviour: training or inference",
+)
 parser.add_argument(
     "--trainImageDir",
     type=str,
@@ -472,7 +478,7 @@ def model_check(
 
 @chris_plugin(
     parser=parser,
-    title="Spleen 3D image segmentation training (MONAI)",
+    title="Spleen 3D image segmentation (MONAI)",
     category="",  # ref. https://chrisstore.co/plugins
     min_memory_limit="16Gi",  # supported units: Mi, Gi
     min_cpu_limit="1000m",  # millicores, e.g. "1000m" = 1 CPU core