BIBSNet: A deep learning baby image brain segmentation network for MRI scans.

1. #!/usr/bin/env python3
2. # coding: utf-8
4. """
5. Wrapper to run nnU-Net_predict with model 512, produced via training on BCP
6. subjects of ages 0-8 months
7. Greg Conan: [email hidden]
8. Created: 2022-02-08
9. Updated: 2022-06-28
10. """
11. # Import standard libraries
12. import argparse
13. from datetime import datetime
14. from glob import glob
15. import os
16. import subprocess
17. import sys
20. def main():
21. # Time how long the script takes and get command-line arguments from user
22. start_time = datetime.now()
23. cli_args = get_cli_args()
25. run_nnUNet_predict(cli_args)
27. # Show user how long the pipeline took and end the pipeline here
28. exit_with_time_info(start_time)
31. def run_nnUNet_predict(cli_args):
32. """
33. Run nnU-Net_predict in a subshell using subprocess
34. :param cli_args: Dictionary containing all command-line input arguments
35. :return: N/A
36. """
37. subprocess.call((cli_args["nnUNet"], "-i",
38. cli_args["input"], "-o", cli_args["output"], "-t",
39. str(cli_args["task"]), "-m", cli_args["model"]))
41. # Only raise an error if there are no output segmentation file(s)
42. if not glob(os.path.join(cli_args["output"], "*.nii.gz")):
43. # NOTE This statement should change if we add a new model
44. sys.exit("Error: Output segmentation file not created at the path "
45. "below during nnUNet_predict run.\n{}\n\nFor your input files "
46. "at the path below, check their filenames and visually "
47. "inspect them if needed.\n{}\n\n"
48. .format(cli_args["output"], cli_args["input"]))
51. def get_cli_args():
52. """
53. :return: Dictionary containing all validated command-line input arguments
54. """
55. script_dir = os.path.dirname(__file__)
56. default_model = "3d_fullres"
57. default_nnUNet_path = os.path.join(script_dir, "nnUNet_predict")
58. default_task_ID = 512
59. parser = argparse.ArgumentParser()
60. parser.add_argument(
61. "--input", "-i", type=valid_readable_dir, required=True,
62. help=("Valid path to existing input directory following valid nnU-Net "
63. "naming conventions (T1w files end with _0000.nii.gz and T2w "
64. "end with _0001.nii.gz). There should be exactly 1 T1w file and "
65. "exactly 1 T2w file in this directory.")
66. )
67. parser.add_argument(
68. "--output", "-o", type=valid_output_dir, required=True,
69. )
70. parser.add_argument(
71. "--nnUNet", "-n", type=valid_readable_file, default=default_nnUNet_path,
72. help=("Valid path to existing executable file to run nnU-Net_predict. "
73. "By default, this script will assume that nnU-Net_predict will "
74. "be in the same directory as this script: {}".format(script_dir))
75. )
76. parser.add_argument( # TODO Does this even need to be an argument, or will it always be the default?
77. "--task", "-t", type=valid_whole_number, default=default_task_ID,
78. help=("Task ID, which should be a 3-digit positive integer starting "
79. "with 5 (e.g. 512).")
80. )
81. parser.add_argument( # TODO Does this even need to be an argument, or will it always be the default?
82. "--model", "-m", default=default_model
83. )
84. return validate_cli_args(vars(parser.parse_args()), parser)
87. def validate_cli_args(cli_args, parser):
88. """
89. :param cli_args: Dictionary containing all command-line input arguments
90. :param parser: argparse.ArgumentParser to raise error if anything's invalid
91. :return: cli_args, but with all input arguments validated
92. """
93. # Verify that there is exactly 1 T1w file and exactly 1 T2w file in the
94. # --input directory
95. err_msg = ("There must be exactly 1 T{0}w file in {1} directory, but the "
96. "number of T{0}w files there currently is {2}")
97. t1or2_path_format = os.path.join(cli_args["input"], "*_000{}.nii.gz")
98. for t1or2 in (1, 2):
99. img_files = glob(t1or2_path_format.format(t1or2 - 1))
100. if len(img_files) != 1:
101. parser.error(err_msg.format(t1or2, cli_args["input"],
102. len(img_files)))
104. # TODO Ensure that task ID is a 3-digit number starting with 5?
106. return cli_args
109. def valid_output_dir(path):
110. """
111. Try to make a folder for new files at path; throw exception if that fails
112. :param path: String which is a valid (not necessarily real) folder path
113. :return: String which is a validated absolute path to real writeable folder
114. """
115. return validate(path, lambda x: os.access(x, os.W_OK),
116. valid_readable_dir, "Cannot create directory at {}",
117. lambda y: os.makedirs(y, exist_ok=True))
120. def valid_readable_dir(path):
121. """
122. :param path: Parameter to check if it represents a valid directory path
123. :return: String representing a valid directory path
124. """
125. return validate(path, os.path.isdir, valid_readable_file,
126. "Cannot read directory at '{}'")
129. def valid_readable_file(path):
130. """
131. Throw exception unless parameter is a valid readable filepath string. Use
132. this, not argparse.FileType("r") which leaves an open file handle.
133. :param path: Parameter to check if it represents a valid filepath
134. :return: String representing a valid filepath
135. """
136. return validate(path, lambda x: os.access(x, os.R_OK),
137. os.path.abspath, "Cannot read file at '{}'")
140. def valid_whole_number(to_validate):
141. """
142. Throw argparse exception unless to_validate is a positive integer
143. :param to_validate: Object to test whether it is a positive integer
144. :return: to_validate if it is a positive integer
145. """
146. return validate(to_validate, lambda x: int(x) >= 0, int,
147. "{} is not a positive integer")
150. def validate(to_validate, is_real, make_valid, err_msg, prepare=None):
151. """
152. Parent/base function used by different type validation functions. Raises an
153. argparse.ArgumentTypeError if the input object is somehow invalid.
154. :param to_validate: String to check if it represents a valid object
155. :param is_real: Function which returns true iff to_validate is real
156. :param make_valid: Function which returns a fully validated object
157. :param err_msg: String to show to user to tell them what is invalid
158. :param prepare: Function to run before validation
159. :return: to_validate, but fully validated
160. """
161. try:
162. if prepare:
163. prepare(to_validate)
164. assert is_real(to_validate)
165. return make_valid(to_validate)
166. except (OSError, TypeError, AssertionError, ValueError,
167. argparse.ArgumentTypeError):
168. raise argparse.ArgumentTypeError(err_msg.format(to_validate))
171. def exit_with_time_info(start_time, exit_code=0):
172. """
173. Terminate the pipeline after displaying a message showing how long it ran
174. :param start_time: datetime.datetime object of when the script started
175. :param exit_code: Int, exit code
176. :return: N/A
177. """
178. print("BIBSnet for this subject took this long to run {}: {}"
179. .format("successfully" if exit_code == 0 else "and then crashed",
180. datetime.now() - start_time))
181. sys.exit(exit_code)
184. if __name__ == "__main__":
185. main()

run.py, under other-open · at the source