def main():
config = load_config()
# make sure those values correspond to the ones used during training
in_channels = config['in_channels']
out_channels = config['out_channels']
# use the same upsampling as during training
interpolate = config['interpolate']
# specify the layer ordering used during training
layer_order = config['layer_order']
# should sigmoid be used as a final activation layer
final_sigmoid = config['final_sigmoid']
init_channel_number = config['init_channel_number']
model = UNet3D(in_channels,
out_channels,
init_channel_number=init_channel_number,
final_sigmoid=final_sigmoid,
interpolate=interpolate,
conv_layer_order=layer_order)
model_path = config['model_path']
logger.info(f'Loading model from {model_path}...')
utils.load_checkpoint(model_path, model)
device = config['device']
model = model.to(device)
logger.info('Loading datasets...')
for test_dataset in get_test_datasets(config):
# run the model prediction on the entire dataset
probability_maps = predict(model, test_dataset, out_channels, device)
# save the resulting probability maps
output_file = _get_output_file(test_dataset)
save_predictions(probability_maps, output_file)
def _load_model(self, final_sigmoid, layer_order):
in_channels = 1
out_channels = 2
# use F.interpolate for upsampling
interpolate = True
return UNet3D(in_channels, out_channels, interpolate, final_sigmoid,
layer_order)
def predict(
model: UNet3D,
brain_loader: BrainLoaders,
out_files,
num_workers,
verbose=True,
):
model.eval()
loaders = brain_loader.test_loader(num_workers=num_workers)
with torch.no_grad():
for loader in loaders: # For every image
epoch_start_time = timeit.default_timer()
result = Sticher(brain_loader.dataset.input_shape,
brain_loader.dataset.slices)
for batch in loader: # For every patch in image
img_patch = batch['image']
file_idx = batch['file_idx'].item()
slice_idx = batch['slice_idx'].item()
if verbose:
logging.info(
f"Predicting image {file_idx} slice id: {slice_idx} "
f"of shape {list(img_patch.size())} starting...")
output = model(img_patch)
if verbose:
logging.info(f"Slice id: {slice_idx} outputted.")
output_patch = BrainDataset.post_process(output)
if verbose:
logging.info(f"Slice id: {slice_idx} post-processed.")
result.update(output_patch, slice_idx)
if verbose:
logging.info(f"Slice id: {slice_idx} saved.")
# Once all the patches are done, save the image
result.save(out_files[file_idx],
brain_loader.dataset.get_nib_file(file_idx))
elapsed = timeit.default_timer() - epoch_start_time
logging.info(f'Predicted in {elapsed} seconds')
def _create_model(final_sigmoid, layer_order):
in_channels = 1
out_channels = 2
# use F.interpolate for upsampling
return UNet3D(in_channels,
out_channels,
final_sigmoid=final_sigmoid,
interpolate=True,
conv_layer_order=layer_order)
def __init__(self,
imgs_dir,
img_postfix='T1',
masks_dir=None,
stack_size=16,
stride=14,
mask_net=False,
verbose=False):
self.img_filenames = listdir(imgs_dir)
tags = [
file[:file.find(img_postfix)]
if file.find(img_postfix) != -1 else file[:file.find('.')]
for file in self.img_filenames if not file.startswith('.')
]
self.img_nibs = [get_nii_files(imgs_dir, tag) for tag in tags]
self.img_files = [img.dataobj for img in self.img_nibs]
self.mask_files = None if masks_dir is None else [
get_nii_files(masks_dir, tag).dataobj for tag in tags
]
self.input_shape = self.img_files[0].shape
patch_shape = (stack_size, self.input_shape[1], self.input_shape[2])
stride_shape = (stride, self.input_shape[1], self.input_shape[2])
self.train_indices = []
logging.info(f'Input shape: {self.input_shape}')
logging.info(f'Patch shape: {patch_shape}')
self.slices = build_slices(self.input_shape,
patch_shape=patch_shape,
stride_shape=stride_shape)
logging.info(
f'Creating dataset from {imgs_dir} and {masks_dir}'
f'\nWith {len(self.img_files)} examples and {len(self.slices)} slices each'
)
self.device = torch.device(
'cuda' if torch.cuda.is_available() else 'cpu')
self.length = len(self.img_files) * len(self.slices)
self.verbose = verbose
if mask_net:
logging.info(f'Loading Pre-process Network from {mask_net}')
pre_net = UNet3D(in_channels=1,
out_channels=2,
final_sigmoid=False,
testing=True)
pre_net.load_state_dict(
torch.load(mask_net, map_location=self.device))
pre_net.to(device=self.device)
logging.info(f'Pre-process Network Loaded')
self.mask_net = pre_net
else:
self.mask_net = None
def _create_model(final_sigmoid, layer_order):
in_channels = 1
out_channels = 2
# use F.interpolate for upsampling and 16 initial feature maps to speed up the tests
return UNet3D(in_channels,
out_channels,
init_channel_number=16,
final_sigmoid=final_sigmoid,
interpolate=True,
conv_layer_order=layer_order)
def _create_model(in_channels,
out_channels,
layer_order,
interpolate=False,
final_sigmoid=True):
return UNet3D(in_channels,
out_channels,
interpolate,
final_sigmoid,
conv_layer_order=layer_order)
def main():
def _get_output_file(dataset):
volume_file, volume_ext = os.path.splitext(dataset.path)
return volume_file + '_probabilities' + volume_ext
parser = argparse.ArgumentParser()
parser.add_argument('--model-path', type=str, help='path to the model')
parser.add_argument('--config-path', type=str,
help='path to the dataset config')
parser.add_argument('--in-channels', default=1, type=int,
help='number of input channels')
parser.add_argument('--out-channels', default=6, type=int,
help='number of output channels')
parser.add_argument('--layer-order', type=str,
help="Conv layer ordering, e.g. 'brc' -> BatchNorm3d+ReLU+Conv3D",
default='brc')
parser.add_argument('--interpolate',
help='use F.interpolate instead of ConvTranspose3d',
action='store_true')
args = parser.parse_args()
# make sure those values correspond to the ones used during training
in_channels = args.in_channels
out_channels = args.out_channels
# use F.interpolate for upsampling
interpolate = args.interpolate
# Conv layer ordering e.g. 'cr' is equivalent to Conv3D+ReLU
conv_layer_order = args.layer_order
model = UNet3D(in_channels, out_channels, interpolate=interpolate,
final_sigmoid=True, conv_layer_order=conv_layer_order)
logger.info(f'Loading model from {args.model_path}...')
utils.load_checkpoint(args.model_path, model)
logger.info('Loading datasets...')
raw_volumes = get_raw_volumes(args.config_path)
if torch.cuda.is_available():
device = torch.device('cuda:0')
else:
logger.warning(
'No CUDA device available. Using CPU for predictions')
device = torch.device('cpu')
model = model.to(device)
for raw_volume in raw_volumes:
probability_maps = predict(model, raw_volume, device)
output_file = _get_output_file(raw_volume)
save_predictions(probability_maps, output_file)
def validate(model: UNet3D, brain_loader: BrainLoaders, loss_fnc=None, is_validation=True, quiet=True):
val_losses = RunningAverage()
val_scores = {fnc: RunningAverage() for fnc in METRICS}
loader = brain_loader.validation_loader() if is_validation else brain_loader.test_loader()
with torch.no_grad():
for batch in loader:
img_batch = batch['image']
mask_batch = batch['mask']
file_idx = batch['file_idx'][0]
patch_validator = PatchValidator(device=brain_loader.device, image_shape=brain_loader.dataset.input_shape)
for slice_no, (img_patch, mask_patch) in enumerate(zip(img_batch, mask_batch)):
img_patch = img_patch.unsqueeze(0)
mask_patch = mask_patch.unsqueeze(0)
# forward pass
output = model(img_patch)
# compute the loss
if loss_fnc is not None:
loss = loss_fnc(output, mask_patch)
val_losses.update(loss.item(), n=1)
# if model contains final_activation layer for normalizing logits apply it, otherwise
# the evaluation metric will be incorrectly computed
if hasattr(model, 'final_activation') and model.final_activation is not None and not model.testing:
output = model.final_activation(output)
slices = brain_loader.dataset.slices
patch_validator.update(output, mask_patch, patch_slice=slices[slice_no % len(slices)])
# When we finish an image, calculate the score for it and update mean
if not quiet:
logging.info(f'Image:{file_idx}:')
for fnc in METRICS:
score = patch_validator.calculate_fnc(fnc)
val_scores[fnc].update(score, n=1)
if not quiet:
logging.info(f"\t\t{fnc+' Score:':<30}{score}")
if loss_fnc is not None:
logging.info(f'Validation: Avg. Loss: {val_losses.avg}.')
scores = '\n\t'.join([f"{fnc+' Score:':<30}{val_scores[fnc].avg}" for fnc in METRICS])
logging.info(f'Evaluation Scores: \n\t{scores}')
return {fnc: v.avg for fnc, v in val_scores.items()}
def main():
logger = get_logger('UNet3DTrainer')
config = load_config()
logger.info(config)
# Create loss criterion
loss_criterion = get_loss_criterion(config)
# Create the model
model = UNet3D(config['in_channels'], config['out_channels'],
final_sigmoid=config['final_sigmoid'],
init_channel_number=config['init_channel_number'],
conv_layer_order=config['layer_order'],
interpolate=config['interpolate'])
model = model.to(config['device'])
# Log the number of learnable parameters
logger.info(f'Number of learnable params {get_number_of_learnable_parameters(model)}')
# Create evaluation metric
eval_criterion = get_evaluation_metric(config)
loaders = get_train_loaders(config)
# Create the optimizer
optimizer = _create_optimizer(config, model)
# Create learning rate adjustment strategy
lr_scheduler = _create_lr_scheduler(config, optimizer)
if config['resume'] is not None:
trainer = UNet3DTrainer.from_checkpoint(config['resume'], model,
optimizer, lr_scheduler, loss_criterion,
eval_criterion, loaders,
logger=logger)
else:
trainer = UNet3DTrainer(model, optimizer, lr_scheduler, loss_criterion, eval_criterion,
config['device'], loaders, config['checkpoint_dir'],
max_num_epochs=config['epochs'],
max_num_iterations=config['iters'],
validate_after_iters=config['validate_after_iters'],
log_after_iters=config['log_after_iters'],
logger=logger)
trainer.fit()
def main():
parser = argparse.ArgumentParser(description='3D U-Net predictions')
parser.add_argument('--model-path', required=True, type=str,
help='path to the model')
parser.add_argument('--in-channels', required=True, type=int,
help='number of input channels')
parser.add_argument('--out-channels', required=True, type=int,
help='number of output channels')
parser.add_argument('--interpolate',
help='use F.interpolate instead of ConvTranspose3d',
action='store_true')
parser.add_argument('--layer-order', type=str,
help="Conv layer ordering, e.g. 'brc' -> BatchNorm3d+ReLU+Conv3D",
default='brc')
args = parser.parse_args()
# make sure those values correspond to the ones used during training
in_channels = args.in_channels
out_channels = args.out_channels
# use F.interpolate for upsampling
interpolate = args.interpolate
layer_order = args.layer_order
model = UNet3D(in_channels, out_channels, interpolate=interpolate,
final_sigmoid=True, conv_layer_order=layer_order)
logger.info(f'Loading model from {args.model_path}...')
utils.load_checkpoint(args.model_path, model)
logger.info('Loading datasets...')
if torch.cuda.is_available():
device = torch.device('cuda:0')
else:
logger.warning(
'No CUDA device available. Using CPU for predictions')
device = torch.device('cpu')
model = model.to(device)
dataset, dataset_size = _get_dataset(1)
probability_maps = predict(model, dataset, dataset_size, device)
output_file = os.path.join(os.path.split(args.model_path)[0],
'probabilities.h5')
save_predictions(probability_maps, output_file)
def main():
config = parse_test_config()
# make sure those values correspond to the ones used during training
in_channels = config.in_channels
out_channels = config.out_channels
# use F.interpolate for upsampling
interpolate = config.interpolate
layer_order = config.layer_order
final_sigmoid = config.final_sigmoid
model = UNet3D(in_channels,
out_channels,
init_channel_number=config.init_channel_number,
final_sigmoid=final_sigmoid,
interpolate=interpolate,
conv_layer_order=layer_order)
logger.info(f'Loading model from {config.model_path}...')
utils.load_checkpoint(config.model_path, model)
logger.info('Loading datasets...')
if torch.cuda.is_available():
device = torch.device('cuda:0')
else:
logger.warning('No CUDA device available. Using CPU for prediction...')
device = torch.device('cpu')
model = model.to(device)
patch = tuple(config.patch)
stride = tuple(config.stride)
for test_path in config.test_path:
# create dataset for a given test file
dataset = HDF5Dataset(test_path,
patch,
stride,
phase='test',
raw_internal_path=config.raw_internal_path)
# run the model prediction on the entire dataset
probability_maps = predict(model, dataset, out_channels, device)
# save the resulting probability maps
output_file = f'{os.path.splitext(test_path)[0]}_probabilities.h5'
save_predictions(probability_maps, output_file)
"%b-%d-%Y_%I-%M-%S_%p") + '_' + args.name + args.output
handlers = [logging.FileHandler(output_filename)]
if not args.quiet:
handlers.append(logging.StreamHandler(sys.stdout))
logging.basicConfig(level=logging.INFO,
format='%(levelname)s: %(message)s',
handlers=handlers)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
logging.info(f'Using device {device}')
input_channels = 1
output_channels = 2
net = UNet3D(in_channels=input_channels,
out_channels=output_channels,
final_sigmoid=False)
logging.info(f'Network:\n'
f'\t{input_channels} input channels\n'
f'\t{output_channels} output channels (classes)\n')
if args.load:
net.load_state_dict(torch.load(args.load, map_location=device))
logging.info(f'Model loaded from {args.load}')
net.to(device=device)
# faster convolutions, but more memory
cudnn.benchmark = True
try:
if args.memory:
def train_net(model: UNet3D,
epochs=5,
learning_rate=0.0002,
val_percent=0.1,
test_percent=0.1,
name='U-Net',
tests=None,
patch_size=16,
testing_memory=False,
mask_model=False):
data_set = BrainDataset(dir_img,
'T1',
dir_mask,
stack_size=patch_size,
mask_net=mask_model)
loader = BrainLoaders(data_set,
ratios=[val_percent, test_percent],
files=[None, tests])
train_loader = loader.train_loader()
val_loader = loader.validation_loader()
test_loader = loader.test_loader()
num_images = data_set.num_files()
log_interval = len(train_loader) if num_images < 10 else len(
data_set.slices) * (num_images // 10)
global_step = 0
logging.info(f'''Starting {name} training:
Epochs: {epochs}
Learning rate: {learning_rate}
Training size: {len(train_loader)} slices
Validation size: {len(val_loader)} images
Testing size: {len(test_loader)} images
Log Interval {log_interval}
''')
optimizer = optim.Adam(model.parameters(),
lr=learning_rate,
weight_decay=0.00001)
losses = []
val_scores = {}
for fnc in METRICS:
val_scores[fnc] = []
for epoch in range(epochs):
epoch_loss = 0
epoch_start_time = timeit.default_timer()
log_start_time = timeit.default_timer()
log_loss = RunningAverage()
for batch in train_loader:
model.train()
img = batch['image']
mask = batch['mask']
masks_pred = model(img)
loss = loss_fnc(masks_pred, mask)
epoch_loss += loss.item()
log_loss.update(loss.item(), n=1)
optimizer.zero_grad()
loss.backward()
optimizer.step()
if testing_memory: # When testing patch sizes only one iteration is enough
return
global_step += 1
if global_step % log_interval == 0:
elapsed = timeit.default_timer() - log_start_time
losses.append(log_loss.avg)
logging.info(
f'I: {global_step}, Avg. Loss: {log_loss.avg} in {elapsed} seconds'
)
log_start_time = timeit.default_timer()
log_loss = RunningAverage()
scores = validate(model, loader, is_validation=True, loss_fnc=loss_fnc)
for fnc in METRICS:
val_scores[fnc].append(scores[fnc])
make_dir(dir_checkpoint)
torch.save(model.state_dict(),
dir_checkpoint + f'{name}_epoch{epoch + 1}.pth')
elapsed = timeit.default_timer() - epoch_start_time
logging.info(
f'Epoch: {epoch + 1} Total Loss: {epoch_loss} in {elapsed} seconds'
)
logging.info(f'Checkpoint {epoch + 1} saved !')
plot_cost(losses, name='Loss', model_name=name + str(epoch) + '_')
for fnc in METRICS:
plot_cost(val_scores[fnc],
name='Validation_' + type(fnc).__name__,
model_name=name + str(epoch) + '_')
logging.info('Starting Testing')
validate(model,
loader,
is_validation=False,
loss_fnc=loss_fnc,
quiet=False)
return parser.parse_args()
if __name__ == '__main__':
args = get_args()
output_filename = args.output + '_' + args.name + datetime.now().strftime("_%b-%d-%Y_%I-%M-%S_%p")
handlers = [logging.FileHandler(output_filename)]
if not args.quiet:
handlers.append(logging.StreamHandler(sys.stdout))
logging.basicConfig(level=logging.INFO,
format='%(levelname)s: %(message)s',
handlers=handlers)
net = UNet3D(in_channels=1, out_channels=2, testing=True, final_sigmoid=False)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
net.to(device=device)
logging.info(f'Using device {device}')
logging.info("Testing model {}".format(args.model))
net.load_state_dict(torch.load(args.model, map_location=device))
logging.info(f"On dataset in {args.input} and {args.label}")
data_set = BrainDataset(args.input, 'T1', args.label)
if args.tests:
data_loader = BrainLoaders(data_set, files=[None, args.tests])
else:
data_loader = BrainLoaders(data_set, ratios=[0, 1])
def main():
parser = argparse.ArgumentParser(description='3D U-Net predictions')
parser.add_argument('--model-path',
required=True,
type=str,
help='path to the model')
parser.add_argument('--in-channels',
required=True,
type=int,
help='number of input channels')
parser.add_argument('--out-channels',
required=True,
type=int,
help='number of output channels')
parser.add_argument('--interpolate',
help='use F.interpolate instead of ConvTranspose3d',
action='store_true')
parser.add_argument(
'--average-channels',
help=
'average the probability_maps across the the channel axis (use only if your channels refer to the same semantic class)',
action='store_true')
parser.add_argument(
'--layer-order',
type=str,
help="Conv layer ordering, e.g. 'crg' -> Conv3D+ReLU+GroupNorm",
default='crg')
parser.add_argument(
'--loss',
type=str,
required=True,
help=
'Loss function used for training. Possible values: [ce, bce, wce, dice]. Has to be provided cause loss determines the final activation of the model.'
)
parser.add_argument('--test-path',
type=str,
required=True,
help='path to the test dataset')
parser.add_argument(
'--patch',
required=True,
type=int,
nargs='+',
default=None,
help='Patch shape for used for prediction on the test set')
parser.add_argument(
'--stride',
required=True,
type=int,
nargs='+',
default=None,
help='Patch stride for used for prediction on the test set')
args = parser.parse_args()
# make sure those values correspond to the ones used during training
in_channels = args.in_channels
out_channels = args.out_channels
# use F.interpolate for upsampling
interpolate = args.interpolate
layer_order = args.layer_order
final_sigmoid = _final_sigmoid(args.loss)
model = UNet3D(in_channels,
out_channels,
final_sigmoid=final_sigmoid,
interpolate=interpolate,
conv_layer_order=layer_order)
logger.info(f'Loading model from {args.model_path}...')
utils.load_checkpoint(args.model_path, model)
logger.info('Loading datasets...')
if torch.cuda.is_available():
device = torch.device('cuda:0')
else:
logger.warning('No CUDA device available. Using CPU for predictions')
device = torch.device('cpu')
model = model.to(device)
patch = tuple(args.patch)
stride = tuple(args.stride)
dataset = HDF5Dataset(args.test_path, patch, stride, phase='test')
probability_maps = predict(model, dataset, dataset.raw.shape, device)
output_file = f'{os.path.splitext(args.test_path)[0]}_probabilities.h5'
save_predictions(probability_maps, output_file, args.average_channels)
def main():
parser = argparse.ArgumentParser(description='3D U-Net predictions')
parser.add_argument('--cdmodel-path', required=True, type=str,
help='path to the coordinate detector model.')
parser.add_argument('--model-path', required=True, type=str,
help='path to the segmentation model')
parser.add_argument('--in-channels', type=int, default=1,
help='number of input channels (default: 1)')
parser.add_argument('--out-channels', type=int, default=2,
help='number of output channels (default: 2)')
parser.add_argument('--init-channel-number', type=int, default=64,
help='Initial number of feature maps in the encoder path which gets doubled on every stage (default: 64)')
parser.add_argument('--layer-order', type=str,
help="Conv layer ordering, e.g. 'crg' -> Conv3D+ReLU+GroupNorm",
default='crg')
parser.add_argument('--final-sigmoid',
action='store_true',
help='if True apply element-wise nn.Sigmoid after the last layer otherwise apply nn.Softmax')
parser.add_argument('--test-path', type=str, nargs='+', required=True, help='path to the test dataset')
parser.add_argument('--raw-internal-path', type=str, default='raw')
parser.add_argument('--patch', type=int, nargs='+', default=None,
help='Patch shape for used for prediction on the test set')
parser.add_argument('--stride', type=int, nargs='+', default=None,
help='Patch stride for used for prediction on the test set')
parser.add_argument('--report-metrics', action='store_true',
help='Whether to print metrics for each prediction')
parser.add_argument('--output-path', type=str, default='./output/',
help='The output path to generate the nifti file')
args = parser.parse_args()
# Check if output path exists
if not os.path.isdir(args.output_path):
os.mkdir(args.output_path)
# make sure those values correspond to the ones used during training
in_channels = args.in_channels
out_channels = args.out_channels
# use F.interpolate for upsampling
interpolate = True
layer_order = args.layer_order
final_sigmoid = args.final_sigmoid
# Define model
UNet_model = UNet3D(in_channels, out_channels,
init_channel_number=args.init_channel_number,
final_sigmoid=final_sigmoid,
interpolate=interpolate,
conv_layer_order=layer_order)
Coor_model = CoorNet(in_channels)
# Define metrics
loss = nn.MSELoss(reduction='sum')
acc = DiceCoefficient()
logger.info('Loading trained coordinate detector model from ' + args.cdmodel_path)
utils.load_checkpoint(args.cdmodel_path, Coor_model)
logger.info('Loading trained segmentation model from ' + args.model_path)
utils.load_checkpoint(args.model_path, UNet_model)
# Load the model to the device
if torch.cuda.is_available():
device = torch.device('cuda:0')
else:
logger.warning('No CUDA device available. Using CPU for predictions')
device = torch.device('cpu')
UNet_model = UNet_model.to(device)
Coor_model = Coor_model.to(device)
# Apply patch training if assigned
if args.patch and args.stride:
patch = tuple(args.patch)
stride = tuple(args.stride)
# Initialise counters
total_dice = 0
total_loss = 0
count = 0
tmp_created = False
for test_path in args.test_path:
if test_path.endswith('.nii.gz'):
if args.report_metrics:
raise ValueError("Cannot report metrics on original files.")
# Temporary save as h5 file
# Preprocess if dim != 192 x 224 x 192
data = preprocess_nifti(test_path, args.output_path)
logger.info('Preprocessing complete.')
hf = h5py.File(test_path + '.h5', 'w')
hf.create_dataset('raw', data=data)
hf.close()
test_path += '.h5'
tmp_created = True
if not args.patch and not args.stride:
curr_shape = np.array(h5py.File(test_path, 'r')[args.raw_internal_path]).shape
patch = curr_shape
stride = curr_shape
# Initialise dataset
dataset = HDF5Dataset(test_path, patch, stride, phase='test', raw_internal_path=args.raw_internal_path)
file_name = test_path.split('/')[-1].split('.')[0]
# Predict the centre coordinates
x, y, z = predict(Coor_model, dataset, out_channels, device)
# Perform segmentation
probability_maps = predict(UNet_model, dataset, out_channels, device, x, y, z)
res = np.argmax(probability_maps, axis=0)
# Put the image batch back to mask with the original size
res = recover_patch(res, x, y, z, dataset.raw.shape)
# Extract LH and RH segmentations and write as file
LH = np.zeros(res.shape)
LH[int(res.shape[0]/2):,:,:] = res[int(res.shape[0]/2):,:,:]
RH = np.zeros(res.shape)
RH[:int(res.shape[0]/2),:,:] = res[:int(res.shape[0]/2),:,:]
LH_img = nib.Nifti1Image(LH, AFF)
RH_img = nib.Nifti1Image(RH, AFF)
nib.save(LH_img, args.output_path + file_name + '_LH.nii.gz')
nib.save(RH_img, args.output_path + file_name + '_RH.nii.gz')
logger.info('File saved to ' + args.output_path + file_name + '_LH.nii.gz')
logger.info('File saved to ' + args.output_path + file_name + '_RH.nii.gz')
if tmp_created:
os.remove(test_path)
if args.report_metrics:
count += 1
# Compute coordinate accuracy
# Coordinate evaluation disabled by default, since not all data have coordinate information
# coor_dataset = HDF5Dataset(test_path, patch, stride, phase='val', raw_internal_path=args.raw_internal_path, label_internal_path='coor')
# coor_target = coor_dataset[0][1].to(device)
# coor_pred_tensor = torch.from_numpy(np.array([x, y, z])).to(device)
# curr_coor_loss = loss(coor_pred_tensor, coor_target)
# total_loss += curr_coor_loss
# logger.info('Current coordinate loss: %f' % (curr_coor_loss))
# Compute segmentation Dice score
label_dataset = HDF5Dataset(test_path, patch, stride, phase='val', raw_internal_path=args.raw_internal_path, label_internal_path='label')
label_target = label_dataset[0][1].to(device)
res_dice = probability_maps
new_shape = np.append(res_dice.shape[0], np.array(label_target.size()))
res_dice = recover_patch_4d(res_dice, x, y, z, new_shape)
pred_tensor = torch.from_numpy(res_dice).to(device).float()
label_target = label_target.view((1,) + label_target.shape)
curr_dice_score = acc(pred_tensor, label_target.long())
total_dice += curr_dice_score
logger.info('Current Dice score: %f' % (curr_dice_score))
# Compute length estimation
logger.info('RH length: ' + str(get_total_dist(res[:int(res.shape[0]/2),:,:])))
logger.info('LH length: ' + str(get_total_dist(res[int(res.shape[0]/2):,:,:])))
if args.report_metrics:
# logger.info('Average loss: %f.' % (total_loss/count))
logger.info('Average Dice score: %f.' % (total_dice/count))
def main():
logger = get_logger('UNet3DTrainer')
# Get device to train on
device = torch.device("cuda:0" if torch.cuda.is_available() else 'cpu')
config = parse_train_config()
logger.info(config)
# Create loss criterion
if config.loss_weight is not None:
loss_weight = torch.tensor(config.loss_weight)
loss_weight = loss_weight.to(device)
else:
loss_weight = None
loss_criterion = get_loss_criterion(config.loss, loss_weight,
config.ignore_index)
model = UNet3D(config.in_channels,
config.out_channels,
init_channel_number=config.init_channel_number,
conv_layer_order=config.layer_order,
interpolate=config.interpolate,
final_sigmoid=config.final_sigmoid)
model = model.to(device)
# Log the number of learnable parameters
logger.info(
f'Number of learnable params {get_number_of_learnable_parameters(model)}'
)
# Create evaluation metric
eval_criterion = get_evaluation_metric(config.eval_metric,
ignore_index=config.ignore_index)
# Get data loaders. If 'bce' or 'dice' loss is used, convert labels to float
train_path, val_path = config.train_path, config.val_path
if config.loss in ['bce']:
label_dtype = 'float32'
else:
label_dtype = 'long'
train_patch = tuple(config.train_patch)
train_stride = tuple(config.train_stride)
val_patch = tuple(config.val_patch)
val_stride = tuple(config.val_stride)
logger.info(f'Train patch/stride: {train_patch}/{train_stride}')
logger.info(f'Val patch/stride: {val_patch}/{val_stride}')
pixel_wise_weight = config.loss == 'pce'
loaders = get_loaders(train_path,
val_path,
label_dtype=label_dtype,
raw_internal_path=config.raw_internal_path,
label_internal_path=config.label_internal_path,
train_patch=train_patch,
train_stride=train_stride,
val_patch=val_patch,
val_stride=val_stride,
transformer=config.transformer,
pixel_wise_weight=pixel_wise_weight,
curriculum_learning=config.curriculum,
ignore_index=config.ignore_index)
# Create the optimizer
optimizer = _create_optimizer(config, model)
if config.resume:
trainer = UNet3DTrainer.from_checkpoint(config.resume,
model,
optimizer,
loss_criterion,
eval_criterion,
loaders,
logger=logger)
else:
trainer = UNet3DTrainer(
model,
optimizer,
loss_criterion,
eval_criterion,
device,
loaders,
config.checkpoint_dir,
max_num_epochs=config.epochs,
max_num_iterations=config.iters,
max_patience=config.patience,
validate_after_iters=config.validate_after_iters,
log_after_iters=config.log_after_iters,
logger=logger)
trainer.fit()
nargs='+',
default=None,
help='Files to use as test cases')
return parser.parse_args()
if __name__ == '__main__':
logging.basicConfig(level=logging.INFO,
format='%(levelname)s: %(message)s')
args = get_args()
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
logging.info(f'Using device {device}')
input_channels = 1
output_channels = 2
net = UNet3D(in_channels=input_channels, out_channels=output_channels)
logging.info(f'Network:\n'
f'\t{input_channels} input channels\n'
f'\t{output_channels} output channels (classes)\n')
if args.load:
net.load_state_dict(torch.load(args.load, map_location=device))
logging.info(f'Model loaded from {args.load}')
net.to(device=device)
# faster convolutions, but more memory
# cudnn.benchmark = True
try:
train_net(model=net,
epochs=args.epochs,
"--verbose",
default=False,
action="store_true",
help="Log more detail")
return parser.parse_args()
if __name__ == "__main__":
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(message)s')
args = get_args()
input_channels = 1
output_channels = 2
net = UNet3D(in_channels=input_channels,
out_channels=output_channels,
testing=True,
final_sigmoid=False)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
net.to(device=device)
logging.info(f'Using device {device}')
base_path = Path(__file__).parent
file_path = (base_path / "../models/test.csv").resolve()
pre_net = False
if args.type == 'mask':
model = (base_path / "../models/BrainMask.pth").resolve()
elif args.type == 'wm':
model = (base_path / "../models/WhiteMatter.pth").resolve()
pre_net = (base_path / "../models/BrainMask.pth").resolve()
required=True)
return parser.parse_args()
if __name__ == "__main__":
logging.basicConfig(level=logging.INFO,
format='%(levelname)s: %(message)s')
args = get_args()
in_files = listdir(args.input)
out_files = [f'OUT_{file}' for file in in_files]
input_channels = 1
output_channels = 2
net = UNet3D(in_channels=input_channels,
out_channels=output_channels,
testing=True)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
net.to(device=device)
logging.info(f'Using device {device}')
logging.info("Loading model {}".format(args.model))
net.load_state_dict(torch.load(args.model, map_location=device))
logging.info("Model loaded !")
for i, fn in enumerate(in_files):
logging.info("\nPredicting image {} ...".format(fn))
nib_img = nib.load(join(args.input, fn))
def train_net(model: UNet3D,
device,
loss_fnc=DiceLoss(sigmoid_normalization=False),
eval_criterion=MeanIoU(),
epochs=5,
batch_size=1,
learning_rate=0.0002,
val_percent=0.04,
test_percent=0.1,
name='U-Net',
save_cp=True,
tests=None):
data_set = BasicDataset(dir_img, dir_mask, 'T1', device)
train_loader, val_loader, test_loader = data_set.split_to_loaders(
val_percent, test_percent, batch_size, test_files=tests)
writer = SummaryWriter(comment=f'LR_{learning_rate}_BS_{batch_size}')
global_step = 0
logging.info(f'''Starting {name} training:
Epochs: {epochs}
Batch size: {batch_size}
Learning rate: {learning_rate}
Training size: {len(train_loader)}
Validation size: {len(val_loader)}
Testing size: {len(test_loader)}
Checkpoints: {save_cp}
Device: {device.type}
''')
optimizer = optim.Adam(model.parameters(),
lr=learning_rate,
weight_decay=0.00001)
losses = []
val_scores = []
for epoch in range(epochs):
epoch_loss = 0
for batch in train_loader:
model.train()
start_time = timeit.default_timer()
img = batch['image']
mask = batch['mask']
masks_pred = model(img)
loss = loss_fnc(masks_pred, mask)
epoch_loss += loss.item()
losses.append(loss.item())
writer.add_scalar('Loss/train', loss.item(), global_step)
optimizer.zero_grad()
loss.backward()
optimizer.step()
global_step += 1
elapsed = timeit.default_timer() - start_time
logging.info(
f'I: {global_step}, Loss: {loss.item()} in {elapsed} seconds')
if global_step % (len(train_loader) // (5 * batch_size)) == 0:
val_score = validate(model, val_loader, loss_fnc,
eval_criterion)
val_scores.append(val_score)
writer.add_scalar('Validation/test', val_score, global_step)
if save_cp:
try:
os.mkdir(dir_checkpoint)
logging.info('Created checkpoint directory')
except OSError:
pass
torch.save(model.state_dict(),
dir_checkpoint + f'{name}_epoch{epoch + 1}.pth')
logging.info(f'Epoch: {epoch + 1} Loss: {epoch_loss}')
logging.info(f'Checkpoint {epoch + 1} saved !')
plot_cost(losses, name='Loss' + str(epoch), model_name=name)
plot_cost(val_scores,
name='Validation' + str(epoch),
model_name=name)
writer.close()
def main():
parser = argparse.ArgumentParser(description='3D U-Net predictions')
parser.add_argument('--model-path',
required=True,
type=str,
help='path to the model')
parser.add_argument('--in-channels',
required=True,
type=int,
help='number of input channels')
parser.add_argument('--out-channels',
required=True,
type=int,
help='number of output channels')
parser.add_argument(
'--init-channel-number',
type=int,
default=64,
help=
'Initial number of feature maps in the encoder path which gets doubled on every stage (default: 64)'
)
parser.add_argument('--interpolate',
help='use F.interpolate instead of ConvTranspose3d',
action='store_true')
parser.add_argument(
'--layer-order',
type=str,
help="Conv layer ordering, e.g. 'crg' -> Conv3D+ReLU+GroupNorm",
default='crg')
parser.add_argument(
'--final-sigmoid',
action='store_true',
help=
'if True apply element-wise nn.Sigmoid after the last layer otherwise apply nn.Softmax'
)
parser.add_argument('--test-path',
type=str,
required=True,
help='path to the test dataset')
parser.add_argument('--raw-internal-path', type=str, default='raw')
parser.add_argument(
'--patch',
required=True,
type=int,
nargs='+',
default=None,
help='Patch shape for used for prediction on the test set')
parser.add_argument(
'--stride',
required=True,
type=int,
nargs='+',
default=None,
help='Patch stride for used for prediction on the test set')
args = parser.parse_args()
# make sure those values correspond to the ones used during training
in_channels = args.in_channels
out_channels = args.out_channels
# use F.interpolate for upsampling
interpolate = args.interpolate
layer_order = args.layer_order
final_sigmoid = args.final_sigmoid
model = UNet3D(in_channels,
out_channels,
init_channel_number=args.init_channel_number,
final_sigmoid=final_sigmoid,
interpolate=interpolate,
conv_layer_order=layer_order)
logger.info(f'Loading model from {args.model_path}...')
utils.load_checkpoint(args.model_path, model)
logger.info('Loading datasets...')
if torch.cuda.is_available():
device = torch.device('cuda:0')
else:
logger.warning('No CUDA device available. Using CPU for predictions')
device = torch.device('cpu')
model = model.to(device)
patch = tuple(args.patch)
stride = tuple(args.stride)
dataset = HDF5Dataset(args.test_path,
patch,
stride,
phase='test',
raw_internal_path=args.raw_internal_path)
probability_maps = predict(model, dataset, out_channels, device)
output_file = f'{os.path.splitext(args.test_path)[0]}_probabilities.h5'
# average channels only in case of final_sigmoid
save_predictions(probability_maps, output_file, final_sigmoid)
def main():
parser = _arg_parser()
logger = get_logger('Trainer')
# Get device to train on
device = torch.device("cuda:0" if torch.cuda.is_available() else 'cpu')
args = parser.parse_args()
if args.loss_weight is not None:
loss_weight = torch.tensor(args.loss_weight)
loss_weight = loss_weight.to(device)
else:
loss_weight = None
if args.network == 'cd':
args.loss = 'mse'
loss_criterion = get_loss_criterion('mse', loss_weight,
args.ignore_index)
model = CoorNet(args.in_channels)
model = model.to(device)
accuracy_criterion = PrecisionBasedAccuracy(30)
elif args.network == 'seg':
if not args.loss:
raise ValueError("Invalid loss assigned.")
loss_criterion = get_loss_criterion(args.loss, loss_weight,
args.ignore_index)
model = UNet3D(args.in_channels,
args.out_channels,
init_channel_number=args.init_channel_number,
conv_layer_order=args.layer_order,
interpolate=True,
final_sigmoid=args.final_sigmoid)
model = model.to(device)
accuracy_criterion = DiceCoefficient(ignore_index=args.ignore_index)
else:
raise ValueError(
"Incorrect network type defined by the --network argument, either cd or seg."
)
# Get data loaders. If 'bce' or 'dice' loss is used, convert labels to float
train_path = args.train_path
if args.loss in ['bce', 'mse']:
label_dtype = 'float32'
else:
label_dtype = 'long'
train_patch = tuple(args.train_patch)
train_stride = tuple(args.train_stride)
pixel_wise_weight = args.loss == 'pce'
loaders = get_loaders(train_path,
label_dtype=label_dtype,
raw_internal_path=args.raw_internal_path,
label_internal_path=args.label_internal_path,
train_patch=train_patch,
train_stride=train_stride,
transformer=args.transformer,
pixel_wise_weight=pixel_wise_weight,
curriculum_learning=args.curriculum,
ignore_index=args.ignore_index)
# Create the optimizer
optimizer = _create_optimizer(args, model)
if args.resume:
trainer = UNet3DTrainer.from_checkpoint(args.resume,
model,
optimizer,
loss_criterion,
accuracy_criterion,
loaders,
logger=logger)
else:
trainer = UNet3DTrainer(model,
optimizer,
loss_criterion,
accuracy_criterion,
device,
loaders,
args.checkpoint_dir,
max_num_epochs=args.epochs,
max_num_iterations=args.iters,
max_patience=args.patience,
validate_after_iters=args.validate_after_iters,
log_after_iters=args.log_after_iters,
logger=logger)
trainer.fit()