import argparse
import shutil
import pickle
import copy
import os
import numpy as np
class BaseArgParser(object):
def __init__(self):
self.parser = argparse.ArgumentParser()
def namespace_to_dict(self, args):
"""Turns a nested Namespace object to a nested dictionary"""
args_dict = vars(copy.deepcopy(args))
for arg in args_dict:
obj = args_dict[arg]
if isinstance(obj, argparse.Namespace):
args_dict[arg] = self.namespace_to_dict(obj)
return args_dict
def fix_nested_namespaces(self, args):
"""Makes sure that nested Namespace work. Supports only one level of nesting."""
group_name_keys = []
for key in args.__dict__:
if '.' in key:
group, name = key.split('.')
group_name_keys.append((group, name, key))
for group, name, key in group_name_keys:
if group not in args:
args.__dict__[group] = argparse.Namespace()
args.__dict__[group].__dict__[name] = args.__dict__[key]
del args.__dict__[key]
def parse_args(self):
args = self.parser.parse_args()
args = self.namespace_to_dict(args)
self.fix_nested_namespaces(args)
return args
class PreproArgParser(BaseArgParser):
def __init__(self):
super(PreproArgParser, self).__init__()
self.parser.add_argument('--in_locs', type=str, required=True,
help='Comma-separated list of paths to all data folders.')
self.parser.add_argument('--modalities', type=str, required=True,
help='Comma-separated list of all input modalities to use.')
self.parser.add_argument('--truth', type=str, required=True,
help='Truth label pattern to use.')
self.parser.add_argument('--create_val', action='store_true', default=False,
help='Whether to create validation set.')
self.parser.add_argument('--out_loc', type=str, default='./data',
help='Location to write preprocessed data.')
def parse_args(self):
args = self.parser.parse_args()
# Create list of all input datasets.
args.in_locs = args.in_locs.split(',')
# Create list of all accepted modalities.
args.modalities = args.modalities.split(',')
# Create output directory if it doesn't already exist.
if os.path.isdir(args.out_loc):
shutil.rmtree(args.out_loc)
args.train_loc = os.path.join(args.out_loc, 'train')
args.val_loc = os.path.join(args.out_loc, 'val')
os.mkdir(args.out_loc)
os.mkdir(args.train_loc)
os.mkdir(args.val_loc)
return args
class TrainArgParser(BaseArgParser):
def __init__(self):
super(TrainArgParser, self).__init__()
# Data args.
self.parser.add_argument('--train_loc', type=str, required=True,
help='Location of .tfrecords training data.')
self.parser.add_argument('--prepro_loc', type=str, required=True,
help='Location of preprocessed dump.')
self.parser.add_argument('--val_loc', type=str, default='',
help='Location of .tfrecords validation data.')
# Checkpoint args.
self.parser.add_argument('--save_folder', type=str, default='',
help='Output folder to save checkpoints, logs, and configs.')
self.parser.add_argument('--load_folder', type=str, default='',
help='Input folder to load checkpoints and configs to resume training.')
# Training args.
self.parser.add_argument('--lr', type=float, default=1e-4,
help='Initial learning rate for training.')
self.parser.add_argument('--batch_size', type=int, default=1,
help='Batch size to use in training.')
self.parser.add_argument('--patience', type=int, default=-1,
help='Number of epochs without validation improvement to stop training.')
self.parser.add_argument('--n_epochs', type=int, default=300,
help='Number of epochs to train for.')
self.parser.add_argument('--gpu', action='store_true', default=False,
help='Whether to train using GPU.')
# Augmentation args.
self.parser.add_argument('--crop_size', type=str, default='128,128,128',
help='Crop size of image (comma-separated h,w,d).')
# Model args.
self.parser.add_argument('--data_format', type=str, dest='model_args.data_format',
default='channels_first', choices=['channels_last', 'channels_first'],
help='Data format to be passed through the model.')
self.parser.add_argument('--base_filters', type=int, dest='model_args.base_filters', default=32,
help='Number of filters in the base convolutional layer.')
self.parser.add_argument('--depth', type=int, dest='model_args.depth', default=4,
help='Number of spatial levels through the model.')
self.parser.add_argument('--l2_scale', type=float, dest='model_args.l2_scale', default=1e-5,
help='Scale of L2 regularization applied to all kernels.')
self.parser.add_argument('--dropout', type=float, dest='model_args.dropout', default=0.2,
help='Dropout ratio to apply to input data.')
self.parser.add_argument('--groups', type=int, dest='model_args.groups', default=8,
help='Number of groups in group normalization.')
self.parser.add_argument('--reduction', type=int, dest='model_args.reduction', default=8,
help='Size of reduction ratio in squeeze-excitation layers.')
self.parser.add_argument('--downsampling', type=str, dest='model_args.downsampling',
default='conv', choices=['conv', 'max', 'avg'],
help='Type of downsampling method.')
self.parser.add_argument('--upsampling', type=str, dest='model_args.upsampling',
default='conv', choices=['conv', 'linear'],
help='Type of upsampling method.')
self.parser.add_argument('--out_ch', type=int, dest='model_args.out_ch', default=3,
help='Number of output classes.')
def parse_args(self):
args = self.parser.parse_args()
# Fix nested Namespaces.
self.fix_nested_namespaces(args)
args.data_format = args.model_args.data_format
# Check data format and GPU compatibility.
args.device = '/device:GPU:0' if args.gpu else '/cpu:0'
if not args.gpu:
assert args.model_args.data_format == 'channels_last', \
'tf.keras.layers.Conv3D only supports `channels_last` input for CPU.'
# Convert model args to dictionaries.
args.model_args = self.namespace_to_dict(args.model_args)
# Set crop size.
args.crop_size = args.crop_size.split(',')
args.crop_size = [int(s) for s in args.crop_size]
# Load preprocessed stats.
prepro = np.load(args.prepro_loc).item()
args.prepro_size = [prepro['size']['h'], prepro['size']['w'], prepro['size']['d'], prepro['size']['c']]
# Check that sizes work out.
assert (args.model_args['base_filters'] / 2) % args.model_args['groups'] == 0, \
'Base filters must be a multiple of {} for group normalization at lowest spatial level.'.format(args.model_args['groups'] * 2)
assert args.model_args['base_filters'] % args.model_args['reduction'] == 0, \
'Base filters must be a multiple of {} for squeeze-excitation reduction.'.format(args.model_args['reduction'])
# Add args.model_args.in_ch for output size of variational autoencoder.
args.model_args['in_ch'] = prepro['size']['c']
# Check for checkpointing option.
if args.load_folder:
with open(os.path.join(args.load_folder, 'train_args.pkl'), 'rb') as f:
chkpt_args = pickle.load(f)
args.model_args = chkpt_args['model_args']
args.crop_size = chkpt_args.crop_size
assert isinstance(args.model_args, dict)
args.save_folder = args.load_folder
# Create checkpoint folder if necessary.
if not os.path.isdir(args.save_folder):
os.mkdir(args.save_folder)
# Save training args.
with open(os.path.join(args.save_folder, 'train_args.pkl'), 'wb') as f:
pickle.dump(self.namespace_to_dict(args), f)
return args
class TestArgParser(BaseArgParser):
def __init__(self):
super(TestArgParser, self).__init__()
# Data.
self.parser.add_argument('--in_locs', type=str, required=True,
help='Comma-separated paths of test data.')
self.parser.add_argument('--modalities', type=str, required=True,
help='Comma-separated modalities to be used as input')
self.parser.add_argument('--truth', type=str, default='',
help='Truth label pattern to use (optional).')
# Training and preprocessing stats.
self.parser.add_argument('--tumor_prepro', type=str, required=True,
help='Path to Numpy preprocessing dump for tumor segmentation.')
self.parser.add_argument('--skull_prepro', type=str, default='',
help='Path to Numpy preprocessing dump for skull segmentation.')
self.parser.add_argument('--tumor_model', type=str, required=True,
help='Path to checkpoint folder for tumor segmentation.')
self.parser.add_argument('--skull_model', type=str, default='',
help='Path to checkpoint folder for skull-stripping segmentation.')
# Input normalization parameters.
self.parser.add_argument('--order', type=int, default=3,
help='Order of interpolation function to be used in voxel resizing.')
self.parser.add_argument('--mode', type=str, default='reflect',
help='Method of handling image edges in interpolation.')
# Test time augmentation and segmentation.
self.parser.add_argument('--spatial_tta', action='store_true', default=True,
help='Whether to apply spatial augmentation on all spatial axes.')
self.parser.add_argument('--channel_tta', type=int, default=0,
help='Additional intensity shifting samples to take.')
self.parser.add_argument('--threshold', type=float, default=0.5,
help='Threshold at which to create mask from probabilities.')
self.parser.add_argument('--gpu', action='store_true', default=False,
help='Whether to evaluate on GPU.')
def parse_args(self):
args = self.parser.parse_args()
args.modalities = args.modalities.split(',')
args.in_locs = args.in_locs.split(',')
# Assert proper combination of inputs.
assert args.threshold > 0 and args.threshold < 1, \
'Threshold must be a probability between (0, 1).'
if args.skull_model:
assert args.skull_prepro, 'Need skull preprocessing stats if model is provided.'
args.skull_strip = bool(args.skull_model)
# Load model args.
with open(os.path.join(args.tumor_model, 'train_args.pkl'), 'rb') as f:
train_args = pickle.load(f)
args.tumor_model_args = train_args['model_args']
args.tumor_spatial_res = 2 ** args.tumor_model_args['depth']
args.tumor_crop_size = train_args['crop_size']
if args.skull_model:
with open(os.path.join(args.skull_model, 'train_args.pkl'), 'rb') as f:
train_args = pickle.load(f)
args.skull_model_args = train_args['model_args']
args.skull_spatial_res = 2 ** args.skull_model_args['depth']
args.skull_crop_size = train_args['crop_size']
# Load prepro stats.
args.tumor_prepro = np.load(args.tumor_prepro).item()
if args.skull_prepro:
args.skull_prepro = np.load(args.skull_prepro).item()
# Check data format and GPU compatibility.
args.device = '/device:GPU:0' if args.gpu else '/cpu:0'
if not args.gpu:
assert args.tumor_model_args['data_format'] == 'channels_last' and args.skull_model_args['data_format'], \
'tf.keras.layers.Conv3D only supports `channels_last` input for CPU.'
return args
