Shepp logan phantom (#203)

Pull from branch shepp-logan-phantom
NKI-AI · May 9, 2022 · d259b5d · d259b5d
1 parent 4e9e99f
commit d259b5d
Show file tree

Hide file tree

Showing 91 changed files with 1,614 additions and 747 deletions.
diff --git a/direct/cli/predict.py b/direct/cli/predict.py
@@ -16,14 +16,16 @@ def register_parser(parser: argparse._SubParsersAction):
         Examples:
         ---------
         Run on single machine:
-            $ direct predict <data_root> <output_directory> --cfg <cfg_path_or_url> --checkpoint <checkpoint_path_or_url>
-                            --num-gpus <num_gpus> [--other-flag-args <other_flags>]
+            $ direct predict <output_directory> --cfg <cfg_path_or_url> --checkpoint <checkpoint_path_or_url> \
+                            --num-gpus <num_gpus> [--data-root <data_root>] [--other-flag-args <other_flags>]
 
         Run on multiple machines:
-            (machine0)$ direct predict <data_root> <output_directory> --cfg <cfg_path_or_url> --checkpoint
-                            <checkpoint_path_or_url> --machine-rank 0 --num-machines 2 --dist-url <URL> [--other-flag-args <other_flags>]
-            (machine1)$ direct predict <data_root> <output_directory> --cfg <cfg_path_or_url> --checkpoint
-                            <checkpoint_path_or_url> --machine-rank 1 --num-machines 2 --dist-url <URL> [--other-flag-args <other_flags>]
+            (machine0)$ direct predict <output_directory> --cfg <cfg_path_or_url> --checkpoint \
+                            <checkpoint_path_or_url> --machine-rank 0 --num-machines 2 [--data-root <data_root>] \
+                                [--dist-url <URL> --other-flag-args <other_flags>]
+            (machine1)$ direct predict <output_directory> --cfg <cfg_path_or_url> --checkpoint \
+                            <checkpoint_path_or_url> --machine-rank 1 --num-machines 2 [--data-root <data_root>] \
+                                [--dist-url <URL> --other-flag-args <other_flags>]
         """
     common_parser = Args(add_help=False)
     predict_parser = parser.add_parser(
@@ -33,8 +35,13 @@ def register_parser(parser: argparse._SubParsersAction):
         epilog=epilog,
         formatter_class=argparse.RawDescriptionHelpFormatter,
     )
-    predict_parser.add_argument("data_root", type=pathlib.Path, help="Path to the inference data directory.")
     predict_parser.add_argument("output_directory", type=pathlib.Path, help="Path to the output directory.")
+    predict_parser.add_argument(
+        "--data-root",
+        type=pathlib.Path,
+        help="Path to the inference data directory (if required by inference dataset).",
+        required=False,
+    )
     predict_parser.add_argument(
         "--cfg",
         dest="cfg_file",

diff --git a/direct/cli/train.py b/direct/cli/train.py
@@ -15,10 +15,10 @@ def register_parser(parser: argparse._SubParsersAction):
         Examples:
         ---------
         Run on single machine:
-            $ direct train training_set validation_set experiment_dir --num-gpus 8 --cfg cfg.yaml
+            $ direct train experiment_dir --num-gpus 8 --cfg cfg.yaml [--training-root training_set --validation-root validation_set]
         Run on multiple machines:
-            (machine0)$ direct train training_set validation_set experiment_dir --machine-rank 0 --num-machines 2 --dist-url <URL> [--other-flags]
-            (machine1)$ direct train training_set validation_set experiment_dir --machine-rank 1 --num-machines 2 --dist-url <URL> [--other-flags]
+            (machine0)$ direct train experiment_dir --machine-rank 0 --num-machines 2 --dist-url <URL> [--training-root training_set --validation-root validation_set] [--other-flags]
+            (machine1)$ direct train experiment_dir --machine-rank 1 --num-machines 2 --dist-url <URL> [--training-root training_set --validation-root validation_set] [--other-flags]
         """
     common_parser = Args(add_help=False)
     train_parser = parser.add_parser(
@@ -28,14 +28,15 @@ def register_parser(parser: argparse._SubParsersAction):
         epilog=epilog,
         formatter_class=argparse.RawDescriptionHelpFormatter,
     )
-
-    train_parser.add_argument("training_root", type=pathlib.Path, help="Path to the training data.")
-    train_parser.add_argument("validation_root", type=pathlib.Path, help="Path to the validation data.")
     train_parser.add_argument(
         "experiment_dir",
         type=pathlib.Path,
         help="Path to the experiment directory.",
     )
+    train_parser.add_argument("--training-root", type=pathlib.Path, help="Path to the training data.", required=False)
+    train_parser.add_argument(
+        "--validation-root", type=pathlib.Path, help="Path to the validation data.", required=False
+    )
     train_parser.add_argument(
         "--cfg",
         dest="cfg_file",

diff --git a/direct/cli/utils.py b/direct/cli/utils.py
@@ -1,9 +1,11 @@
 # coding=utf-8
 # Copyright (c) DIRECT Contributors
+
 import argparse
 import pathlib
 import sys
 
+from direct.types import FileOrUrl, PathOrString
 from direct.utils.io import check_is_valid_url
 
 
@@ -14,13 +16,15 @@ def is_file(path):
     raise argparse.ArgumentTypeError(f"{path} is not a valid file or url.")
 
 
-def file_or_url(path):
+def file_or_url(path: PathOrString) -> FileOrUrl:
     if check_is_valid_url(path):
-        return path
-    path = pathlib.Path(path)
-    if path.is_file():
-        return path
-    raise argparse.ArgumentTypeError(f"{path} is not a valid file or url.")
+        return FileOrUrl(path)
+    else:
+        path = pathlib.Path(path)
+        if path.is_file():
+            return FileOrUrl(path)
+        else:
+            raise argparse.ArgumentTypeError(f"{path} is not a valid file or url.")
 
 
 def check_train_val(key, name):