/
cli.py
117 lines (100 loc) · 3.29 KB
/
cli.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
pyrobex.cli
console script for pyrobex
this is a silly script because it provides a python script
interface to a bash script, but it's slightly easier to
deploy because of setup.py and it allows easy testing of
the pyrobex package
Author: Jacob Reinhold (jcreinhold@gmail.com)
Created on: May 6, 2021
"""
import argparse
import logging
import sys
from typing import List, Optional, Union
ArgType = Optional[Union[argparse.Namespace, List[str]]]
def setup_log(verbosity: int) -> None:
"""get logger with appropriate logging level and message"""
if verbosity == 1:
level = logging.getLevelName("INFO")
elif verbosity >= 2:
level = logging.getLevelName("DEBUG")
else:
level = logging.getLevelName("WARNING")
logging.basicConfig(
format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", level=level
)
logging.captureWarnings(True)
def arg_parser() -> argparse.ArgumentParser:
desc = "ROBEX v1.2 T1-w Brain Extraction"
parser = argparse.ArgumentParser(description=desc)
required = parser.add_argument_group("Required")
required.add_argument(
"t1_image", type=str, help="path to T1-w image to be skull-stripped"
)
options = parser.add_argument_group("Options")
options.add_argument(
"-os",
"--output-stripped",
type=str,
default=None,
help="path to output stripped image",
)
options.add_argument(
"-om",
"--output-mask",
type=str,
default=None,
help="path to output brain mask",
)
options.add_argument(
"-s",
"--seed",
type=int,
default=0,
help="random seed for reproducible results",
)
options.add_argument(
"-v",
"--verbosity",
action="count",
default=0,
help="increase output verbosity (e.g., -vv is more than -v)",
)
return parser
def check_args(args: argparse.Namespace) -> None:
if args.output_stripped is None and args.output_mask is None:
msg = "--output-stripped and/or --output-mask should be specified,\n"
msg += "otherwise this script has no output.\nAborting."
from pyrobex.errors import PyRobexError
raise PyRobexError(msg)
def main(args: ArgType = None) -> int:
"""Console script for pyrobex."""
parser = arg_parser()
if args is None:
args = parser.parse_args()
elif isinstance(args, list):
args = parser.parse_args(args)
setup_log(args.verbosity)
logger = logging.getLogger(__name__)
# import pyrobex here for logging backend
from pyrobex.io import NiftiImage
from pyrobex.robex import robex
check_args(args)
image = NiftiImage.load(args.t1_image)
stripped, mask = robex(image, args.seed)
if args.output_stripped is not None:
stripped.to_filename(args.output_stripped)
msg = f"Output skull-stripped image saved to: {args.output_stripped}."
logger.info(msg)
if args.output_mask is not None:
mask.to_filename(args.output_mask)
msg = f"Output mask image saved to: {args.output_mask}"
logger.info(msg)
return 0
if __name__ == "__main__":
parser = arg_parser()
args = parser.parse_args()
sys.exit(main(args)) # pragma: no cover