#!/usr/bin/env python3
"""
Use ``vstats_hemi_to_avg`` (``h2a``) from UNRAVEL to automatically average hemisphere images with their mirrored counterparts. This can also smooth the images with a kernel and apply a mask.
Inputs:
- input_img_LH.nii.gz
- input_img_RH.nii.gz
Outputs:
- input_img_LRavg.nii.gz
- input_img_s100_LRavg.nii.gz
Usage:
------
vstats_hemi_to_avg [--kernel 0.1] [--axis 0] [--shift 2] [--parallel] [--atlas_mask path/atlas_mask.nii.gz] [-v]
"""
import numpy as np
import nibabel as nib
from pathlib import Path
from rich import print
from rich.traceback import install
from concurrent.futures import ThreadPoolExecutor
from fsl.wrappers import fslmaths
from unravel.core.help_formatter import RichArgumentParser, SuppressMetavar, SM
from unravel.core.config import Configuration
from unravel.core.utils import log_command, match_files, verbose_start_msg, verbose_end_msg, print_func_name_args_times
from unravel.voxel_stats.apply_mask import load_mask
from unravel.voxel_stats.mirror import mirror
[docs]
def parse_args():
parser = RichArgumentParser(formatter_class=SuppressMetavar, add_help=False, docstring=__doc__)
opts = parser.add_argument_group('Optional arguments')
opts.add_argument('-k', '--kernel', help='Smoothing kernel radius in mm if > 0. Default: 0 ', default=0, type=float, action=SM)
opts.add_argument('-ax', '--axis', help='Axis to flip the image along. Default: 0', default=0, type=int, action=SM)
opts.add_argument('-s', '--shift', help='Number of voxels to shift content after flipping. Default: 2', default=2, type=int, action=SM)
opts.add_argument('-tp', '--parallel', help='Enable parallel processing with thread pools', default=False, action='store_true')
opts.add_argument('-amas', '--atlas_mask', help='path/atlas_mask.nii.gz', default=None, action=SM)
general = parser.add_argument_group('General arguments')
general.add_argument('-v', '--verbose', help='Increase verbosity. Default: False', default=False, action='store_true')
return parser.parse_args()
[docs]
@print_func_name_args_times()
def hemi_to_LR_avg(lh_file, rh_file, kernel=0, axis=0, shift=2, atlas_mask=None):
path = lh_file.parent
output_filename = rh_file.name.replace('_RH.nii.gz', f'_s{str(int(kernel * 1000))}_LRavg.nii.gz' if kernel > 0 else '_LRavg.nii.gz')
output_path = path / output_filename
# Check if the output file already exists
if output_path.exists():
print(f"Output {output_filename} already exists. Skipping...")
return
# Load images
right_nii = nib.load(str(rh_file))
left_nii = nib.load(str(lh_file))
# Optionally smooth images
if kernel > 0:
print(f" Smoothing images with a kernel radius of {kernel} mm")
right_nii = fslmaths(right_nii).s(kernel).run()
left_nii = fslmaths(left_nii).s(kernel).run()
right_img = np.asanyarray(right_nii.dataobj, dtype=right_nii.header.get_data_dtype()).squeeze()
left_img = np.asanyarray(left_nii.dataobj, dtype=left_nii.header.get_data_dtype()).squeeze()
# Mirror and average images
mirrored_left_img = mirror(left_img, axis=axis, shift=shift)
averaged_img = (right_img + mirrored_left_img) / 2
# Apply the mask
if atlas_mask is not None:
mask_img = load_mask(atlas_mask)
averaged_img[~mask_img] = 0 # Use logical NOT to flip True/False
# Save the averaged image
averaged_nii = nib.Nifti1Image(averaged_img, right_nii.affine, right_nii.header)
nib.save(averaged_nii, output_path)
print(f" Saved averaged image to {output_filename}")
[docs]
@log_command
def main():
install()
args = parse_args()
Configuration.verbose = args.verbose
verbose_start_msg()
path = Path.cwd()
rh_files = match_files('*_RH.nii.gz', path)
if args.parallel:
with ThreadPoolExecutor() as executor:
executor.map(lambda rh_file: hemi_to_LR_avg(path / str(rh_file).replace('_RH.nii.gz', '_LH.nii.gz'), rh_file, args.kernel, args.axis, args.shift, args.atlas_mask), rh_files)
else:
for rh_file in rh_files:
lh_file = path / str(rh_file).replace('_RH.nii.gz', '_LH.nii.gz')
hemi_to_LR_avg(lh_file, rh_file, args.kernel, args.axis, args.shift, args.atlas_mask)
verbose_end_msg()
if __name__ == '__main__':
main()