forked from CBICA/BrainMaGe
-
Notifications
You must be signed in to change notification settings - Fork 0
/
brain_mage_run
executable file
·165 lines (145 loc) · 8.49 KB
/
brain_mage_run
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
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
#!usr/bin/env python
# -*- coding: utf-8 -*-
"""
Created on Sat May 30 01:05:59 2020
@author: siddhesh
"""
from __future__ import absolute_import, print_function, division
import argparse
import os
import pandas as pd
from BrainMaGe.trainer import trainer_main
from BrainMaGe.tester import test_ma, test_multi_4
import pkg_resources
if __name__ == '__main__':
parser = argparse.ArgumentParser(prog='BrainMaGe', formatter_class=argparse.RawTextHelpFormatter,
description='\nThis code was implemented for Deep Learning '+\
'based training and inference of 3D-U-Net,\n3D-Res-U-Net models for '+\
'Brain Extraction a.k.a Skull Stripping in biomedical NIfTI volumes.\n'+\
'The project is hosted at: https://github.com/CBICA/BrainMaGe * \n'+\
'See the documentation for details on its use.\n'+\
'If you are using this tool, please cite out paper.'
'This software accompanies the research presented in:\n'+\
'Thakur et al., \'Brain Extraction on MRI Scans in Presence of Diffuse\n'+\
'Glioma:Multi-institutional Performance Evaluation of Deep Learning Methods'+\
'and Robust Modality-Agnostic Training\'.\n'+\
'DOI: 10.1016/j.neuroimage.2020.117081\n' +\
'We hope our work helps you in your endeavours.\n'+ '\n'\
'Copyright: Center for Biomedical Image Computing and Analytics (CBICA), University of Pennsylvania.\n'\
'For questions and feedback contact: [email protected]')
parser.add_argument('-params', dest='params', type=str,
help='Specify the architecture of the model to be used, by providing a\n'+\
'config file [PARAMS_CFG]. A sample of the files is stored in\n'+\
'BrainMaGe/config folder for the train, test. Checkout the parameter\n'+\
'explanation in the Readme.md for more details.\n',
required=True)
parser.add_argument('-train', dest='train', type=str,
help='Should be set to "True" (without the quotes) if you are trying to\n'+\
'run training, but make sure you intensity standardize the data \n'+\
'before attempting to train.\n', default='False')
parser.add_argument('-test', dest='test', type=str,
help='Should be set to "False" (without the quotes) if you are trying\n'+\
'to train a new model, do not set the training to true as testing\n'+\
'will be overridden.\n', default='True')
parser.add_argument('-dev', default='0', dest='device', type=str,
help='used to set on which device the prediction will run.\n'+
'Must be either int or str. Use int for GPU id or\n'+
'\'cpu\' to run on CPU. Avoid training on CPU. \n'+
'Default for selecting first GPU is set to -dev 0\n',
required=False)
parser.add_argument('-mode', dest='mode', type=str,
help='Should be one of "MA" or "Multi-4" without the quotes so that \n'+
'the appropriate weight files are loaded automatically during\n'+\
'the test time.')
parser.add_argument('-save_brain', default=1, type=int, required=False, dest='save_brain',
help='if set to 0 the segmentation mask will be only produced and\n'+\
'and the mask will not be applied on the input image to produce\n'+\
' a brain. This step is to be only applied if you trust this\n'+\
'software and do not feel the need for Manual QC. This will save\n'+\
' you some time. This is useless for training though.')
parser.add_argument('-load', default=None, dest='load', type=str,
help='If the location of the weight file is passed, the internal methods\n'+\
'are overridden to apply these weights to the model. We warn against\n'+\
'the usage of this unless you know what you are passing. C')
parser.add_argument('-v', '--version', action='version',
version=pkg_resources.require("BrainMaGe")[0].version + '\n\nCopyright: Center for Biomedical Image Computing and Analytics (CBICA), University of Pennsylvania.', help="Show program's version number and exit.")
args = parser.parse_args()
params_file = os.path.abspath(args.params)
DEVICE = args.device
# Reading in all the parameters
mode = args.mode
save_brain = args.save_brain
# some sanity checking
if args.train == args.test:
raise ValueError('Please enable either testing or training modes, not both')
if args.train == False and args.test == False:
raise ValueError('One of the options needs to be enabled.')
# If weights are given in params, then set weights to given params
# else set weights to None
if args.load is not None:
weights = os.path.abspath(args.load)
else:
weights = None
# If weights are not None, which meeans the weights are given
# Then check if weights are .ckpt for training
# and .pt for testing
# Else raise value error
if weights is not None:
if os.path.exists(weights):
if args.train == 'True':
_, ext = os.path.splitext(weights)
if ext != '.ckpt':
raise ValueError("The extension was not a .ckpt file for training to enable proper\n"+\
"resume during training. Please pass a .ckpt file.")
elif args.test == 'True':
print(args.mode)
if args.mode.lower() == 'ma' or args.mode.lower() == 'multi_4' or args.mode.lower() == 'bids':
_, ext = os.path.splitext(weights)
if ext != '.pt':
raise ValueError("Expected a .pt file, got a file with %s extension. If it is a\n"+\
".ckpt file, please conver it with our converion script\n"+\
"mentioned in the Readme.md")
else :
raise ValueError('Unknown value for mode. Expected one of "MA" or "Multi-4" without the quotes.',
'We received : ', args.mode,
'Common mistakes include spelling mistakes, check it to make sure.')
else:
if args.train == 'True':
pass
elif args.test == 'True':
base_dir = os.path.dirname(os.path.dirname(os.path.realpath(__file__)))
base_dir = os.path.join(os.path.dirname(base_dir), 'BrainMaGe/weights')
if args.mode.lower() == 'ma' or args.mode.lower() == 'bids':
weights = os.path.join(base_dir, 'resunet_ma.pt')
elif args.mode.lower() == 'multi-4':
weights = os.path.join(base_dir, 'resunet_multi_4.pt')
else:
raise ValueError('Unknown value for mode. Expected one of "MA" or "Multi-4" without the quotes.',
'We received : ', args.mode,
'Common mistakes include spelling mistakes, check it to make sure.')
print("Weight file used :", weights)
print(__file__)
if DEVICE == 'cpu':
pass
else:
DEVICE = int(DEVICE)
if args.save_brain == 0:
args.save_brain = False
elif args.save_brain == 1:
args.save_brain = True
else:
raise ValueError('Unknown value for save brain:')
if args.train == 'True':
trainer_main.train_network(params_file, DEVICE, weights)
elif args.test == 'True':
if args.mode.lower() == 'ma' or args.mode.lower() == 'bids':
test_ma.infer_ma(params_file, DEVICE, args.save_brain, weights)
elif args.mode.lower() == 'multi-4':
test_multi_4.infer_multi_4(params_file, DEVICE, args.save_brain, weights)
else :
raise ValueError('Unknown value for mode. Expected one of "MA" or "Multi-4" without the quotes.',
'We received : ', args.mode,
'Common mistakes include spelling mistakes, check it to make sure.')
else:
raise ValueError("Expected the modes to be set with either -train True or -test True.\n"+\
"Please try again!")