split_to_lesions.py

from collections import defaultdict
import os.path
import nibabel
import numpy as np
from pirec.artefacts import NiiGzImage
from skimage.measure import label as sklabel


def make_lesional_masks(masks):
    """Separate lesions to individual masks with 4D connectivity."""
    unique_lesion_id = 0
    seen = defaultdict(set)
    new_masks = []

    nib_masks = [nibabel.load(m.filename) for m in masks]
    l_masks = [sklabel(m.get_data()) for m in nib_masks]
    for orig_mask_i in range(len(l_masks)):
        n_lesions = l_masks[orig_mask_i].max()
        next_mask_i = orig_mask_i + 1
        for bl_lesion_i in range(1, n_lesions + 1):
            if bl_lesion_i in seen[id(l_masks[orig_mask_i])]:
                # skip lesion if it was 4D-linked to a previous timepoint
                continue
            union = l_masks[orig_mask_i] == bl_lesion_i
            for fup_mask in l_masks[next_mask_i:len(l_masks)]:
                n_fup_lesions = fup_mask.max()
                for fup_lesion_i in range(1, n_fup_lesions + 1):
                    this_fup_lesion = fup_mask == fup_lesion_i
                    if np.any(np.logical_and(union, this_fup_lesion)):
                        union = np.logical_or(union, this_fup_lesion)
                        seen[id(fup_mask)].add(fup_lesion_i)
            out = nibabel.Nifti1Image(union.astype(np.uint8), nib_masks[orig_mask_i].affine)
            filename = '{0}_lesion{1:03d}.nii.gz'.format(
                masks[orig_mask_i].justname,
                unique_lesion_id
            )
            full_path = os.path.join(masks[orig_mask_i].dirname, filename)
            nibabel.save(out, full_path)
            unique_lesion_id += 1
            new_masks.append(NiiGzImage(full_path))
    return new_masks