utils.py

import importlib
import io
import logging
import os
import shutil
import sys
import uuid

import h5py
import matplotlib.pyplot as plt
import numpy as np
import scipy.sparse as sparse
import torch
from PIL import Image
# from sklearn.decomposition import PCA

plt.ioff()
plt.switch_backend('agg')


def save_checkpoint(state, is_best, checkpoint_dir, logger=None):
    """Saves model and training parameters at '{checkpoint_dir}/last_checkpoint.pytorch'.
    If is_best==True saves '{checkpoint_dir}/best_checkpoint.pytorch' as well.

    Args:
        state (dict): contains model's state_dict, optimizer's state_dict, epoch
            and best evaluation metric value so far
        is_best (bool): if True state contains the best model seen so far
        checkpoint_dir (string): directory where the checkpoint are to be saved
    """

    def log_info(message):
        if logger is not None:
            logger.info(message)

    if not os.path.exists(checkpoint_dir):
        log_info(
            f"Checkpoint directory does not exists. Creating {checkpoint_dir}")
        os.mkdir(checkpoint_dir)

    last_file_path = os.path.join(checkpoint_dir, 'last_checkpoint.pytorch')
    log_info(f"Saving last checkpoint to '{last_file_path}'")
    torch.save(state, last_file_path)
    if is_best:
        best_file_path = os.path.join(checkpoint_dir, 'best_checkpoint.pytorch')
        log_info(f"Saving best checkpoint to '{best_file_path}'")
        shutil.copyfile(last_file_path, best_file_path)


def load_checkpoint(checkpoint_path, model, optimizer=None):
    """Loads model and training parameters from a given checkpoint_path
    If optimizer is provided, loads optimizer's state_dict of as well.

    Args:
        checkpoint_path (string): path to the checkpoint to be loaded
        model (torch.nn.Module): model into which the parameters are to be copied
        optimizer (torch.optim.Optimizer) optional: optimizer instance into
            which the parameters are to be copied

    Returns:
        state
    """
    if not os.path.exists(checkpoint_path):
        raise IOError(f"Checkpoint '{checkpoint_path}' does not exist")

    state = torch.load(checkpoint_path)
    model.load_state_dict(state['model_state_dict'])

    if optimizer is not None:
        optimizer.load_state_dict(state['optimizer_state_dict'])

    return state


def save_network_output(output_path, output, logger=None):
    if logger is not None:
        logger.info(f'Saving network output to: {output_path}...')
    output = output.detach().cpu()[0]
    with h5py.File(output_path, 'w') as f:
        f.create_dataset('predictions', data=output, compression='gzip')


def get_logger(name, level=logging.INFO):
    logger = logging.getLogger(name)
    logger.setLevel(level)
    # Logging to console
    stream_handler = logging.StreamHandler(sys.stdout)
    formatter = logging.Formatter(
        '%(asctime)s [%(threadName)s] %(levelname)s %(name)s - %(message)s')
    stream_handler.setFormatter(formatter)
    logger.addHandler(stream_handler)

    return logger


def get_number_of_learnable_parameters(model):
    model_parameters = filter(lambda p: p.requires_grad, model.parameters())
    return sum([np.prod(p.size()) for p in model_parameters])


class RunningAverage:
    """Computes and stores the average
    """

    def __init__(self):
        self.count = 0
        self.sum = 0
        self.avg = 0

    def update(self, value, n=1):
        self.count += n
        self.sum += value * n
        self.avg = self.sum / self.count


def find_maximum_patch_size(model, device):
    """Tries to find the biggest patch size that can be send to GPU for inference
    without throwing CUDA out of memory"""
    logger = get_logger('PatchFinder')
    in_channels = model.in_channels

    patch_shapes = [(64, 128, 128), (96, 128, 128),
                    (64, 160, 160), (96, 160, 160),
                    (64, 192, 192), (96, 192, 192)]

    for shape in patch_shapes:
        # generate random patch of a given size
        patch = np.random.randn(*shape).astype('float32')

        patch = torch \
            .from_numpy(patch) \
            .view((1, in_channels) + patch.shape) \
            .to(device)

        logger.info(f"Current patch size: {shape}")
        model(patch)


def unpad(patch, index, shape, pad_width=4):
    """
    Remove `pad_width` voxels around the edges of a given patch.
    """

    def _new_slices(slicing, max_size):
        if slicing.start == 0:
            p_start = 0
            i_start = 0
        else:
            p_start = pad_width
            i_start = slicing.start + pad_width

        if slicing.stop == max_size:
            p_stop = None
            i_stop = max_size
        else:
            p_stop = -pad_width
            i_stop = slicing.stop - pad_width

        return slice(p_start, p_stop), slice(i_start, i_stop)

    D, H, W = shape

    i_c, i_z, i_y, i_x = index
    p_c = slice(0, patch.shape[0])

    p_z, i_z = _new_slices(i_z, D)
    p_y, i_y = _new_slices(i_y, H)
    p_x, i_x = _new_slices(i_x, W)

    patch_index = (p_c, p_z, p_y, p_x)
    index = (i_c, i_z, i_y, i_x)
    return patch[patch_index], index


def create_feature_maps(init_channel_number, number_of_fmaps):
    return [init_channel_number * 2 ** k for k in range(number_of_fmaps)]


# Code taken from https://github.com/cremi/cremi_python
def adapted_rand(seg, gt, all_stats=False):
    """Compute Adapted Rand error as defined by the SNEMI3D contest [1]
    Formula is given as 1 - the maximal F-score of the Rand index
    (excluding the zero component of the original labels). Adapted
    from the SNEMI3D MATLAB script, hence the strange style.
    Parameters
    ----------
    seg : np.ndarray
        the segmentation to score, where each value is the label at that point
    gt : np.ndarray, same shape as seg
        the groundtruth to score against, where each value is a label
    all_stats : boolean, optional
        whether to also return precision and recall as a 3-tuple with rand_error
    Returns
    -------
    are : float
        The adapted Rand error; equal to $1 - \frac{2pr}{p + r}$,
        where $p$ and $r$ are the precision and recall described below.
    prec : float, optional
        The adapted Rand precision. (Only returned when `all_stats` is ``True``.)
    rec : float, optional
        The adapted Rand recall.  (Only returned when `all_stats` is ``True``.)
    References
    ----------
    [1]: http://brainiac2.mit.edu/SNEMI3D/evaluation
    """
    # just to prevent division by 0
    epsilon = 1e-6

    # segA is truth, segB is query
    segA = np.ravel(gt)
    segB = np.ravel(seg)
    n = segA.size

    n_labels_A = np.amax(segA) + 1
    n_labels_B = np.amax(segB) + 1

    ones_data = np.ones(n)

    p_ij = sparse.csr_matrix((ones_data, (segA[:], segB[:])), shape=(n_labels_A, n_labels_B))

    a = p_ij[1:n_labels_A, :]
    b = p_ij[1:n_labels_A, 1:n_labels_B]
    c = p_ij[1:n_labels_A, 0].todense()
    d = b.multiply(b)

    a_i = np.array(a.sum(1))
    b_i = np.array(b.sum(0))

    sumA = np.sum(a_i * a_i)
    sumB = np.sum(b_i * b_i) + (np.sum(c) / n)
    sumAB = np.sum(d) + (np.sum(c) / n)

    precision = sumAB / max(sumB, epsilon)
    recall = sumAB / max(sumA, epsilon)

    fScore = 2.0 * precision * recall / max(precision + recall, epsilon)
    are = 1.0 - fScore

    if all_stats:
        return are, precision, recall
    else:
        return are


class _TensorboardFormatter:
    """
    Tensorboard formatters converts a given batch of images (be it input/output to the network or the target segmentation
    image) to a series of images that can be displayed in tensorboard. This is the parent class for all tensorboard
    formatters which ensures that returned images are in the 'CHW' format.
    """

    def __init__(self, **kwargs):
        pass

    def __call__(self, name, batch):
        """
        Transform a batch to a series of tuples of the form (tag, img), where `tag` corresponds to the image tag
        and `img` is the image itself.

        Args:
             name (str): one of 'inputs'/'targets'/'predictions'
             batch (torch.tensor): 4D or 5D torch tensor
        """

        def _check_img(tag_img):
            tag, img = tag_img

            assert img.ndim == 2 or img.ndim == 3, 'Only 2D (HW) and 3D (CHW) images are accepted for display'

            if img.ndim == 2:
                img = np.expand_dims(img, axis=0)
            else:
                C = img.shape[0]
                assert C == 1 or C == 3, 'Only (1, H, W) or (3, H, W) images are supported'

            return tag, img

        assert name in ['inputs', 'targets', 'predictions']

        tagged_images = self.process_batch(name, batch)

        return list(map(_check_img, tagged_images))

    def process_batch(self, name, batch):
        raise NotImplementedError


class DefaultTensorboardFormatter(_TensorboardFormatter):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)

    def process_batch(self, name, batch):
        tag_template = '{}/batch_{}/channel_{}/slice_{}'

        tagged_images = []

        if batch.ndim == 5:
            # NCDHW
            slice_idx = batch.shape[2] // 2  # get the middle slice
            for batch_idx in range(batch.shape[0]):
                for channel_idx in range(batch.shape[1]):
                    tag = tag_template.format(name, batch_idx, channel_idx, slice_idx)
                    img = batch[batch_idx, channel_idx, slice_idx, ...]
                    tagged_images.append((tag, self._normalize_img(img)))
        else:
            # batch hafrom sklearn.decomposition import PCAs no channel dim: NDHW
            slice_idx = batch.shape[1] // 2  # get the middle slice
            for batch_idx in range(batch.shape[0]):
                tag = tag_template.format(name, batch_idx, 0, slice_idx)
                img = batch[batch_idx, slice_idx, ...]
                tagged_images.append((tag, self._normalize_img(img)))

        return tagged_images

    @staticmethod
    def _normalize_img(img):
        return np.nan_to_num((img - np.min(img)) / np.ptp(img))


class EmbeddingsTensorboardFormatter(DefaultTensorboardFormatter):
    def __init__(self, plot_variance=False, **kwargs):
        super().__init__(**kwargs)
        self.plot_variance = plot_variance

    def process_batch(self, name, batch):
        if name == 'inputs':
            assert batch.ndim == 5
            # skip coordinate channels and take only the first 'raw' channel
            batch = batch[:, 0, ...]
            return super().process_batch(name, batch)
        elif name == 'predictions':
            return self._embeddings_to_rgb(batch)
        else:
            return super().process_batch(name, batch)

    def _embeddings_to_rgb(self, batch):
        assert batch.ndim == 5

        tag_template = 'embeddings/batch_{}/slice_{}'
        tagged_images = []

        slice_idx = batch.shape[2] // 2  # get the middle slice
        for batch_idx in range(batch.shape[0]):
            tag = tag_template.format(batch_idx, slice_idx)
            img = batch[batch_idx, :, slice_idx, ...]  # CHW
            rgb_img = self._pca_project(img)
            tagged_images.append((tag, rgb_img))
            if self.plot_variance:
                cum_explained_variance_img = self._plot_cum_explained_variance(img)
                tagged_images.append((f'cumulative_explained_variance/batch_{batch_idx}', cum_explained_variance_img))

        return tagged_images

    # def _pca_project(self, embeddings):
    #     assert embeddings.ndim == 3
    #     # reshape (C, H, W) -> (C, H * W) and transpose
    #     flattened_embeddings = embeddings.reshape(embeddings.shape[0], -1).transpose()
    #     # init PCA with 3 principal components: one for each RGB channel
    #     pca = PCA(n_components=3)
    #     # fit the model with embeddings and apply the dimensionality reduction
    #     flattened_embeddings = pca.fit_transform(flattened_embeddings)
    #     # reshape back to original
    #     shape = list(embeddings.shape)
    #     shape[0] = 3
    #     img = flattened_embeddings.transpose().reshape(shape)
    #     # normalize to [0, 255]
    #     img = 255 * (img - np.min(img)) / np.ptp(img)
    #     return img.astype('uint8')

    # def _plot_cum_explained_variance(self, embeddings):
    #     # reshape (C, H, W) -> (C, H * W) and transpose
    #     flattened_embeddings = embeddings.reshape(embeddings.shape[0], -1).transpose()
    #     # fit PCA to the data
    #     pca = PCA().fit(flattened_embeddings)
    #
    #     plt.figure()
    #     # plot cumulative explained variance ratio
    #     plt.plot(np.cumsum(pca.explained_variance_ratio_))
    #     plt.xlabel('number of components')
    #     plt.ylabel('cumulative explained variance');
    #     buf = io.BytesIO()
    #     plt.savefig(buf, format='jpeg')
    #     buf.seek(0)
    #     img = np.asarray(Image.open(buf)).transpose(2, 0, 1)
    #     plt.close('all')
    #     return img


def get_tensorboard_formatter(config):
    if config is None:
        return DefaultTensorboardFormatter()

    class_name = config['name']
    m = importlib.import_module('unet3d.utils')
    clazz = getattr(m, class_name)
    return clazz(**config)


def expand_as_one_hot(input, C, ignore_index=None):
    """
    Converts NxDxHxW label image to NxCxDxHxW, where each label gets converted to its corresponding one-hot vector
    :param input: 4D input image (NxDxHxW)
    :param C: number of channels/labels
    :param ignore_index: ignore index to be kept during the expansion
    :return: 5D output image (NxCxDxHxW)
    """
    assert input.dim() == 4

    # expand the input tensor to Nx1xDxHxW before scattering
    input = input.unsqueeze(1)
    # create result tensor shape (NxCxDxHxW)
    shape = list(input.size())
    shape[1] = C

    if ignore_index is not None:
        # create ignore_index mask for the result
        mask = input.expand(shape) == ignore_index
        # clone the src_res tensor and zero out ignore_index in the input
        input = input.clone()
        input[input == ignore_index] = 0
        # scatter to get the one-hot tensor
        result = torch.zeros(shape).to(input.device).scatter_(1, input, 1)
        # bring back the ignore_index in the result
        result[mask] = ignore_index
        return result
    else:
        # scatter to get the one-hot tensor
        return torch.zeros(shape).to(input.device).scatter_(1, input, 1)


def plot_segm(segm, ground_truth, plots_dir='.'):
    """
    Saves predicted and ground truth segmentation into a PNG files (one per channel).

    :param segm: 4D ndarray (CDHW)
    :param ground_truth: 4D ndarray (CDHW)
    :param plots_dir: directory where to save the plots
    """
    assert segm.ndim == 4
    if ground_truth.ndim == 3:
        stacked = [ground_truth for _ in range(segm.shape[0])]
        ground_truth = np.stack(stacked)

    assert ground_truth.ndim == 4

    f, axarr = plt.subplots(1, 2)

    for seg, gt in zip(segm, ground_truth):
        mid_z = seg.shape[0] // 2

        axarr[0].imshow(seg[mid_z], cmap='prism')
        axarr[0].set_title('Predicted segmentation')

        axarr[1].imshow(gt[mid_z], cmap='prism')
        axarr[1].set_title('Ground truth segmentation')

        file_name = f'segm_{str(uuid.uuid4())[:8]}.png'
        plt.savefig(os.path.join(plots_dir, file_name))