These python modules provide functions for dealing with multi label volumetric images (such as anatomical atlases) in .nii format or any similar format compatible with nibabel. Both can also be used as command line tools. mlv_merge.py takes a directory of single-label or multilabel volumetric image files and creates a new nifiti multi-label imag…

mlv_merge.py

#!/usr/bin/env python

import os
import sys
import argparse
import collections
import numpy
import nibabel as nib


def merge_ml_volumes(input_dir, label_join_char="+", voxel_threshold=15):
    shape = None
    available_values = set(range(1, 32000))
    labels = {}
    data = None
    valid_extensions = [".nii", ".nii.gz"]

    niftis = [item for item in os.listdir(input_dir) if any(item.endswith(ext) for ext in valid_extensions)]
    if not niftis:
        if __name__ == "__main__":
            print 'Warning: No nifti image files were found in the input directory! \nTerminating early...'
            sys.exit(1)
        else:
            raise Exception('No nifti image files were found in the input directory!')

    for indx, item in enumerate(niftis):
        image_name = item.split(".")[0]
        new_img = nib.load(input_dir + '/' + item)
        new_data = new_img.get_data().flatten()
        new_values_map = {}

        if shape is None:
            shape = new_img.shape
            data = numpy.zeros(len(new_data)).astype('int16', copy=False)
        if shape != new_img.shape:
            raise Exception('Cannot merge volumes with different dimensions!')

        # Merge image
        for i in numpy.where(new_data != 0)[0]:
            combo = (data[i], new_data[i])
            if not combo in new_values_map:
                try:
                    new_values_map[combo] = available_values.pop()
                except KeyError:
                    raise Exception('Ran out of label values, use the --ml option specify a larger number of label values.')
            data[i] = new_values_map[combo]
        available_values -= set(new_values_map.values())

        # Merge labels
        is_mutlilabel = len([x[0] for x in new_values_map.keys() if x[0] == 0]) > 1
        if len(new_values_map) == 1:
            if new_values_map.keys()[0][0] == 0:
                labels[new_values_map.values()[0]] = image_name
            else:
                labels[new_values_map.values()[0]] = labels[new_values_map.keys()[0][0]] + label_join_char + image_name
        else:
            i = 0
            for combo in new_values_map.keys():
                if combo[0] == 0:
                    if is_mutlilabel:
                        i += 1
                        labels[new_values_map[combo]] = image_name + '_' + str(i)
                    else:
                        labels[new_values_map[combo]] = image_name
                else:
                    combined_label = labels[combo[0]] + label_join_char + image_name
                    if combined_label in labels.values():
                        combined_label += '_' + str(combo[1])
                    labels[new_values_map[combo]] = combined_label

        # Empty out any very small labels resulting from overlap, if this is the last item
        if indx == len(niftis)-1:
            fd = collections.Counter(data)
            for val in fd:
                if fd[val] >= voxel_threshold or not label_join_char in labels[val]:
                    continue
                original_label = label_join_char.join(labels[val].split(label_join_char)[0:-1])
                for v, l in labels.iteritems():
                    if l == original_label:
                        data[data == val] = v
                        break

        # Remove any emptied labels
        emptied_values = set(labels.keys())-set(data)
        for val in emptied_values:
            del labels[val]
            available_values = set([val]+list(available_values))
        sys.stdout.write('.')
        sys.stdout.flush()

    # Fill empty values
    new_vals = {}
    labels2 = {}
    for i, v in enumerate(sorted(labels.keys())):
        new_vals[v] = i+1
        if new_vals[v] != v:
            data[data == v] = i+1
    for v in new_vals:
        labels2[new_vals[v]] = labels[v]
    labels = labels2

    return (nib.Nifti1Image(data.reshape(shape), None), labels)


if __name__ == "__main__":
    label_join_char = "+"
    voxel_threshold = 15

    parser = argparse.ArgumentParser()
    parser.add_argument("input_dir")
    parser.add_argument("--jc", help="specify a character (which should not be in any of the filenames) for joining the names of overlapping labels, defaults to" + label_join_char)
    parser.add_argument("--vt", help="specify the minimum number of voxels required for a label overlap to be propagated, defaults to " + str(voxel_threshold))
    parser.add_argument("--o",  help="specify a valid path for the resulting merged image, defaults to the directory above input_directory")
    args = parser.parse_args()

    # Interpret arguments
    input_dir = args.input_dir
    while input_dir[-1] == "/":
        input_dir = input_dir[0:-1]
    if args.jc:
        label_join_char = args.jc
    if args.vt:
        voxel_threshold = int(args.vt)
    output_dir = args.o or "/".join(input_dir.split("/")[0:-1])
    if os.path.isdir(output_dir):
        output_path = output_dir + '/merged_image.nii.gz'
    elif os.path.isdir("/".join(input_dir.split("/")[0:-1])):
        output_path = output_dir
    else:
        raise Exception('Output directory appears to be invalid!')

    new_image, labels = merge_ml_volumes(input_dir, label_join_char, voxel_threshold)

    # Write new nifit image
    nib.save(new_image, output_path)

    # Write text file with labels
    open(output_dir+"/labels.txt", "w").write("")
    with open(output_dir+"/labels.txt", "a") as labels_file:
        for val in labels:
            labels_file.write(str(val) + " "*(5-len(str(val))) + labels[val] + "\n")
    sys.exit(0)

mlv_surface.py

#!/usr/bin/env python

import os
import argparse
import nibabel as nib


def build_label_surfaces(nifti_image, output_dir):
    if not output_dir[-1] == "/":
        output_dir += "/"
    img = nifti_image.get_data().flatten()
    dims = nifti_image.shape
    index_of = lambda x, y, z: x + y*dims[0] + z*dims[0]*dims[1]
    meshes = set()
    present = None

    for z in xrange(dims[2]):
        for y in xrange(dims[1]):
            for x in xrange(dims[0]):
                previous = present or img[index_of(x - 1, y, z)]
                present = img[index_of(x, y, z)]
                above = img[index_of(x, y-1, z)]
                behind = img[index_of(x, y, z-1)]
                squares = []

                if x > 0 and present != previous:
                        squares.append({
                            "labels": (previous, present),
                            "t1": [[x, y, z], [x, y+1, z], [x, y+1, z+1]],
                            "t2": [[x, y+1, z+1], [x, y, z+1], [x, y, z]],
                            "normal": [1, 0, 0]
                        })
                if y > 0 and present != above:
                        squares.append({
                            "labels": (above, present),
                            "t1": [[x, y, z+1], [x+1, y, z+1], [x+1, y, z]],
                            "t2": [[x+1, y, z], [x, y, z], [x, y, z+1]],
                            "normal": [0, 1, 0]
                        })
                if z > 0 and present != behind:
                        squares.append({
                            "labels": (behind, present),
                            "t1": [[x, y, z], [x+1, y, z], [x+1, y+1, z]],
                            "t2": [[x+1, y+1, z], [x, y+1, z], [x, y, z]],
                            "normal": [0, 0, 1]
                        })

                for square in squares:
                    mesh_id = "-".join(map(str, sorted(square["labels"])))
                    file_path = "{0}{1}.stl".format(output_dir, mesh_id)
                    if not mesh_id in meshes:
                        open(file_path, "w").write("solid voxel_surface")
                        meshes.add(mesh_id)
                    if square["labels"][0] < square["labels"][1]:
                        square["normal"] = [-x for x in square["normal"]]
                        square["t1"][1], square["t1"][2] = square["t1"][2], square["t1"][1]
                        square["t2"][1], square["t2"][2] = square["t2"][2], square["t2"][1]
                    write_square(square["t1"], square["t2"], square["normal"], open(file_path, "a"))

    for mesh_id in meshes:
        file_path = "{0}{1}.stl".format(output_dir, mesh_id)
        open(file_path, "a").write("nendsolid\n")


def write_square(t1, t2, normal, output_file):
    for line in [
        " facet normal {0} {1} {2}".format(*normal),
        "  outer loop",
        "   vertex {0} {1} {2}".format(*t1[0]),
        "   vertex {0} {1} {2}".format(*t1[1]),
        "   vertex {0} {1} {2}".format(*t1[2]),
        "  endloop",
        " endfacet",
        " facet normal {0} {1} {2}".format(*normal),
        "  outer loop",
        "   vertex {0} {1} {2}".format(*t2[0]),
        "   vertex {0} {1} {2}".format(*t2[1]),
        "   vertex {0} {1} {2}".format(*t2[2]),
        "  endloop",
        " endfacet"
    ]:
        output_file.write(line + "\n")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("nift_image")
    parser.add_argument("output_dir")
    args = parser.parse_args()

    nift_image = args.nift_image
    if not os.path.isfile(nift_image):
        raise Exception('Input image file appears to be invalid!')
    output_dir = args.output_dir
    while output_dir[-1] == "/":
        output_dir = output_dir[0:-1]
    if not os.path.isdir(output_dir):
        raise Exception('Output directory appears to be invalid!')

    build_label_surfaces(nib.load(nift_image), output_dir)