import os
import glob
import scipy
import numpy as np
import nibabel as nib
import tensorflow as tf
from tqdm import tqdm
from scipy.ndimage import zoom
from args import TestArgParser
from util import DiceCoefficient
from model import Model
class Interpolator(object):
def __init__(self, modalities, order=3, mode='reflect'):
self.modalities = modalities
self.order = order
self.mode = mode
def __call__(self, path):
"""Extracts Numpy image and normalizes it to 1 mm^3."""
# Extract raw images from each.
image = []
pixdim = []
affine = []
for name in self.modalities:
file_card = glob.glob(os.path.join(path, '*' + name + '*' + '.nii' + '*'))[0]
img = nib.load(file_card)
image.append(np.array(img.dataobj).astype(np.float32))
pixdim.append(img.header['pixdim'][:4])
affine.append(np.stack([img.header['srow_x'],
img.header['srow_y'],
img.header['srow_z'],
np.array([0., 0., 0., 1.])], axis=0))
# Prepare image.
image = np.stack(image, axis=-1)
self.pixdim = np.mean(pixdim, axis=0, dtype=np.float32)
self.affine = np.mean(affine, axis=0, dtype=np.float32)
# Rescale and interpolate voxels spatially.
if np.any(self.pixdim[:-1] != 1.0):
image = zoom(image, self.pixdim[:-1] + [1.0], order=self.order, mode=self.mode)
# Rescale and interpolate voxels depthwise (along time).
if self.pixdim[-1] != 1.0:
image = zoom(image, [1.0, 1.0, self.pixdim[-1], 1.0], order=self.order, mode=self.mode)
# Mask out background voxels.
mask = np.max(image, axis=-1, keepdims=True)
mask = (mask > 0).astype(np.float32)
return image, mask
def reverse(self, output, path):
"""Reverses the interpolation performed in __call__."""
# Scale back spatial voxel interpolation.
if np.any(self.pixdim[:-1] != 1.0):
output = zoom(output, 1.0 / self.pixdim[:-1], order=self.order, mode=self.mode)
# Scale back depthwise voxel interpolation.
if self.pixdim[-1] != 1.0:
output = zoom(output, [1.0, 1.0, 1.0 / self.pixdim[-1]], order=self.order, mode=self.mode)
# Save file.
nib.save(nib.Nifti1Image(output, self.affine),
os.path.join(path, 'mask.nii'))
return output
class TestTimeAugmentor(object):
"""Handles full inference on input with test-time augmentation."""
def __init__(self,
mean,
std,
model,
model_data_format,
spatial_tta=True,
channel_tta=0,
threshold=0.5):
self.mean = mean
self.std = std
self.model = model
self.model_data_format = model_data_format
self.channel_tta = channel_tta
self.threshold = threshold
self.channel_axis = -1 if self.model_data_format == 'channels_last' else 1
self.spatial_axes = [1, 2, 3] if self.model_data_format == 'channels_last' else [2, 3, 4]
if spatial_tta:
self.augment_axes = [self.spatial_axes, []]
for axis in self.spatial_axes:
pairs = self.spatial_axes.copy()
pairs.remove(axis)
self.augment_axes.append([axis])
self.augment_axes.append(pairs)
else:
self.augment_axes = [[]]
def __call__(self, x, bmask):
# Normalize and prepare input (assumes input data format of 'channels_last').
x = (x - self.mean) / self.std
# Transpose to channels_first data format if required by model.
if self.model_data_format == 'channels_first':
x = tf.transpose(x, (3, 0, 1, 2))
bmask = tf.transpose(bmask, (3, 0, 1, 2))
x = tf.expand_dims(x, axis=0)
bmask = tf.expand_dims(bmask, axis=0)
# Initialize list of inputs to feed model.
y = []
# Create shape for intensity shifting.
shape = [1, 1, 1]
shape.insert(self.channel_axis, x.shape[self.channel_axis])
if self.channel_tta:
_, var = tf.nn.moments(x, axes=self.spatial_axes, keepdims=True)
std = tf.sqrt(var)
# Apply spatial augmentation.
for flip in self.augment_axes:
# Run inference on spatially augmented input.
aug = tf.reverse(x, axis=flip)
aug, *_ = self.model(aug, training=False, inference=True)
y.append(tf.reverse(aug, axis=flip))
for _ in range(self.channel_tta):
shift = tf.random.uniform(shape, -0.1, 0.1)
scale = tf.random.uniform(shape, 0.9, 1.1)
# Run inference on channel augmented input.
aug = (aug + shift * std) * scale
aug = self.model(aug, training=False, inference=True)
aug = tf.reverse(aug, axis=flip)
y.append(aug)
# Aggregate outputs.
y = tf.concat(y, axis=0)
y = tf.reduce_mean(y, axis=0, keepdims=True)
# Mask out zero-valued voxels.
y *= bmask
# Take the argmax to determine label.
# y = tf.argmax(y, axis=self.channel_axis, output_type=tf.int32)
# Transpose back to channels_last data format.
y = tf.squeeze(y, axis=0)
if self.model_data_format == 'channels_first':
y = tf.transpose(y, (1, 2, 3, 0))
return y
def pad_to_spatial_res(res, x, mask):
# Assumes that x and mask are channels_last data format.
res = tf.convert_to_tensor([res])
shape = tf.convert_to_tensor(x.shape[:-1], dtype=tf.int32)
shape = res - (shape % res)
pad = [[0, shape[0]],
[0, shape[1]],
[0, shape[2]],
[0, 0]]
orig_shape = list(x.shape[:-1])
x = tf.pad(x, pad, mode='CONSTANT', constant_values=0.0)
mask = tf.pad(mask, pad, mode='CONSTANT', constant_values=0.0)
return x, mask, orig_shape
def main(args):
in_ch = len(args.modalities)
# Initialize model(s) and load weights / preprocessing stats.
tumor_model = Model(**args.tumor_model_args)
tumor_crop_size = args.tumor_model_args['crop_size']
_ = tumor_model(tf.zeros(shape=[1] + tumor_crop_size + [in_ch] if args.tumor_model_args['data_format'] == 'channels_last' \
else [1, in_ch] + tumor_crop_size, dtype=tf.float32))
tumor_model.load_weights(os.path.join(args.tumor_model, 'chkpt.hdf5'))
tumor_mean = tf.convert_to_tensor(args.tumor_prepro['norm']['mean'], dtype=tf.float32)
tumor_std = tf.convert_to_tensor(args.tumor_prepro['norm']['std'], dtype=tf.float32)
if args.skull_strip:
skull_model = Model(**args.skull_model_args)
skull_crop_size = args.skull_model_args['crop_size']
_ = model(tf.zeros(shape=[1] + skull_crop_size + [in_ch] if args.skull_model_args['data_format'] == 'channels_last' \
else [1, in_ch] + skull_crop_size, dtype=tf.float32))
skull_model.load_weights(os.path.join(args.skull_model, 'chkpt.hdf5'))
skull_mean = tf.convert_to_tensor(args.skull_prepro['norm']['mean'], dtype=tf.float32)
skull_std = tf.convert_to_tensor(args.skull_prepro['norm']['std'], dtype=tf.float32)
# Initialize helper classes for inference and evaluation (optional).
dice_fn = DiceCoefficient(data_format='channels_last')
interpolator = Interpolator(args.modalities, order=args.order, mode=args.mode)
tumor_ttaugmentor = TestTimeAugmentor(
tumor_mean,
tumor_std,
tumor_model,
args.tumor_model_args['data_format'],
spatial_tta=args.spatial_tta,
channel_tta=args.channel_tta,
threshold=args.threshold)
if args.skull_strip:
skull_ttaugmentor = TestTimeAugmentor(
skull_mean,
skull_std,
skull_model,
args.skull_model_args['data_format'],
spatial_tta=args.spatial_tta,
channel_tta=args.channel_tta,
threshold=args.threshold)
for loc in args.in_locs:
for path in tqdm(glob.glob(os.path.join(loc, '*'))):
with tf.device(args.device):
# If data is labeled, extract label.
try:
file_card = glob.glob(os.path.join(path, '*' + args.truth + '*' + '.nii' + '*'))[0]
y = np.array(nib.load(file_card).dataobj).astype(np.float32)
y = tf.expand_dims(y, axis=-1)
except:
y = None
# Rescale and interpolate input image.
x, mask = interpolator(path)
# Strip MRI brain of skull and eye sockets.
if args.skull_strip:
x, pad_mask, pad = pad_to_spatial_res( # Pad to spatial resolution
args.skull_spatial_res,
x,
mask)
skull_mask = skull_ttaugmentor(x, pad_mask) # Inference with test time augmentation.
skull_mask = 1.0 - skull_mask # Convert skull positives into negatives.
x *= skull_mask # Mask out skull voxels.
x = tf.slice(x, # Remove padding.
[0, 0, 0, 0],
pad + [-1])
# Label brain tumor categories per voxel.
x, pad_mask, pad = pad_to_spatial_res( # Pad to spatial resolution.
args.tumor_spatial_res,
x,
mask)
tumor_mask = tumor_ttaugmentor(x, pad_mask) # Inference with test time augmentation.
tumor_mask = tf.slice(tumor_mask, # Remove padding.
[0, 0, 0, 0],
pad + [-1])
tumor_mask += 1 # Convert [0,1,2] to [1,2,3] for label consistency.
tumor_mask = tumor_mask.numpy()
np.place(tumor_mask, tumor_mask >= 3, [4]) # Replace label `3` with `4` for label consistency.
# Reverse interpolation and save as .nii.
y_pred = interpolator.reverse(tumor_mask, path)
# If label is available, score the prediction.
if y is not None:
macro, micro = dice_fn(y, y_pred)
print('{}. Macro: {ma: 1.4f}. Micro: {mi: 1.4f}'
.format(path.split('/')[-1], ma=macro, mi=micro), flush=True)
if __name__ == '__main__':
parser = TestArgParser()
args = parser.parse_args()
print('Test args: {}'.format(args))
main(args)
