test.py

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)