MRI Image Mosaic generator

mosaic.py

#!/usr/bin/env python

# Copyright 2024 Board of Regents of the Universities of Wisconsin
# Written by Nate Vack <[email protected]>

"""MRI Image Mosaic.

Makes an MxN mosaic of images, intended for checking things like defacing 
and alignment.

I know there's probably something that does this in one of the image packages
but afni is a disaster and slicer puts everything in one giant column

Usage:
  mosaic.py [options] <input> <output>

Options:
  -h --help         Show this message
  -v --verbose      Print more debugging information
  --grid=<shape>    Shape of the grid, as MxN [default: 1x1]
  --plane=<num>     Which plane you want -- 0, 1, 2 [default: 0]
  --rotations=<num> Number of times you want to rotate 90 degrees [default: 0]
"""

import nibabel as nib
import matplotlib.pyplot as plt
import numpy as np
from docopt import docopt

import logging

logging.basicConfig(format="%(message)s")
logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)

plt.switch_backend("Agg")


def main(infile, outfile, grid_dims, plane, rotations):
    nii = nib.load(infile)
    logger.info(f"Read {infile}")
    data = nii.get_fdata()
    grid_xy = [int(x) for x in grid_dims.split("x")]
    needed_slices = grid_xy[0] * grid_xy[1]
    plane = int(plane or 0)
    rotations = int(rotations or 0)

    # We want to start in the middle of the brain for our slices, so figure out a
    # start and a step
    slices_in_plane = data.shape[plane]
    start = slices_in_plane // (needed_slices + 1)
    step = slices_in_plane // needed_slices

    slice_indexes = np.arange(start, slices_in_plane, step)[:needed_slices]

    # Compute
    voxel_sizes = np.abs(np.diag(nii.affine))[:3]
    other_sizes = voxel_sizes[np.delete(np.arange(3), plane)]
    ratio = other_sizes[1] / other_sizes[0]
    logger.debug(f"Aspect ratio: {ratio}")

    _fig, axes = plt.subplots(grid_xy[0], grid_xy[1], squeeze=False)
    # This will let us find the data we want -- we'll change the value of
    # slicer[plane] to our desired slice index
    slicer = [slice(None)] * 3
    for i, ax in enumerate(axes.flat):
        logger.debug(f"Taking slice {slice_indexes[i]}")
        slicer[plane] = slice_indexes[i]
        slice_data = data[tuple(slicer)]
        logger.debug("Taken")
        rotated = np.rot90(slice_data, k=rotations)
        logger.debug("Rotated")
        ax.imshow(rotated, cmap="gray", aspect=ratio)
        logger.debug("shown")
        ax.axis("off")

    plt.tight_layout(pad=0.1)

    logger.info(f"Wrote {outfile}")
    plt.savefig(outfile, bbox_inches="tight", pad_inches=0, dpi=150)


if __name__ == "__main__":
    args = docopt(__doc__)
    if args["--verbose"]:
        logger.setLevel(logging.DEBUG)
    logger.debug(args)
    main(
        args["<input>"],
        args["<output>"],
        args["--grid"],
        args["--plane"],
        args["--rotations"],
    )