Source code for unravel.image_tools.dilate

#!/usr/bin/env python3

"""
Use ``img_dilate`` (``dilate``) from UNRAVEL to dilate a mask in a NIfTI image.

All positive input voxels are treated as foreground. The output is a binary image.

Usage:
------
    img_dilate -i path/image.nii.gz -it 2 [-o path/image_dil2.nii.gz]
"""

import numpy as np
from rich.traceback import install
from scipy.ndimage import binary_dilation

from unravel.core.img_io import load_3D_img, save_as_nii
from unravel.core.help_formatter import RichArgumentParser, SuppressMetavar, SM
from unravel.core.utils import log_command, verbose_start_msg, verbose_end_msg




def parse_args():
    parser = RichArgumentParser(formatter_class=SuppressMetavar, add_help=False, docstring=__doc__)

    reqs = parser.add_argument_group('Required arguments')
    reqs.add_argument('-i', '--input', help='Path to a NIfTI mask. Non-binary inputs are permitted, but all values > 0 will be treated as foreground.', required=True, action=SM)
    reqs.add_argument('-it', '--iterations', help='Number of dilation iterations. Default: 1', type=int, default=1, action=SM)

    opts = parser.add_argument_group('Optional arguments')
    opts.add_argument('-o', '--output', help='path/image.nii.gz. Default: None (saves as path/image_dilated.nii.gz)', action=SM)

    return parser.parse_args()





@log_command
def main():
    install()
    args = parse_args()
    verbose_start_msg()

    # Load original intensities so non-binary input can be detected
    img = load_3D_img(args.input)

    if np.any((img != 0) & (img != 1)):
        print(
            "\n    [yellow bold]Warning:[/yellow bold] Input image is not binary. "
            "All values > 0 will be treated as foreground.\n"
        )

    img_bin = img > 0
    img_dil = binary_dilation(img_bin, iterations=args.iterations)

    # Save the modified image as a NIfTI file
    output_path = args.output if args.output else str(args.input).replace('.nii.gz', f'_dil{args.iterations}.nii.gz')
    save_as_nii(img_dil, output_path, reference=args.input)

    verbose_end_msg()



if __name__ == '__main__':
    main()