def lsq_normalize(img_dir, mask_dir=None, output_dir=None, write_to_disk=True):
"""
normalize intensities of a set of MR images by minimizing the squared distance
between CSF, GM, and WM means within the set
Args:
img_dir (str): directory containing MR images
mask_dir (str): directory containing masks for MR images
output_dir (str): directory to save images if you do not want them saved in
same directory as data_dir
write_to_disk (bool): write the normalized data to disk or nah
Returns:
normalized (np.ndarray): last normalized image from img_dir
"""
input_files = io.glob_nii(img_dir)
if output_dir is None:
out_fns = [None] * len(input_files)
else:
out_fns = []
for fn in input_files:
_, base, ext = io.split_filename(fn)
out_fns.append(os.path.join(output_dir, base + '_lsq' + ext))
if not os.path.exists(output_dir):
os.mkdir(output_dir)
mask_files = [None] * len(
input_files) if mask_dir is None else io.glob_nii(mask_dir)
standard_tissue_means = None
normalized = None
for i, (img_fn, mask_fn,
out_fn) in enumerate(zip(input_files, mask_files, out_fns)):
_, base, _ = io.split_filename(img_fn)
logger.info(
'Transforming image {} to standard scale ({:d}/{:d})'.format(
base, i + 1, len(input_files)))
img = io.open_nii(img_fn)
mask = io.open_nii(mask_fn) if mask_fn is not None else None
tissue_mem = mask_util.fcm_class_mask(img, mask)
if standard_tissue_means is None:
csf_tissue_mask = find_tissue_mask(img, mask, tissue_type='csf')
csf_normed_data = fcm_normalize(img, csf_tissue_mask).get_fdata()
standard_tissue_means = calc_tissue_means(csf_normed_data,
tissue_mem)
del csf_tissue_mask, csf_normed_data
img_data = img.get_fdata()
tissue_means = calc_tissue_means(img_data, tissue_mem)
sf = find_scaling_factor(tissue_means, standard_tissue_means)
logger.debug('Scaling factor for {}: {:0.3e}'.format(base, sf))
normalized = nib.Nifti1Image(sf * img_data, img.affine, img.header)
if write_to_disk:
io.save_nii(normalized, out_fn, is_nii=True)
return normalized
def csf_mask(img,
brain_mask,
contrast='t1',
csf_thresh=0.9,
return_prob=False,
mrf=0.25,
use_fcm=False):
"""
create a binary mask of csf using atropos (FMM) segmentation
of a T1-w image
Args:
img (ants.core.ants_image.ANTsImage or nibabel.nifti1.Nifti1Image): target img
brain_mask (ants.core.ants_image.ANTsImage or nibabel.nifti1.Nifti1Image): brain mask for img
contrast (str): contrast of the img (e.g., t1, t2, or flair)
csf_thresh (float): membership threshold to count as CSF
return_prob (bool): if true, then return membership values
instead of binary (i.e., thresholded membership) mask
mrf (float): markov random field parameter
(i.e., smoothness parameter, higher is a smoother segmentation)
use_fcm (bool): use FCM segmentation instead of atropos (may be less accurate)
cannot use return_prob flag
Returns:
csf (np.ndarray): binary CSF mask for img
"""
# convert nibabel to antspy format images (to do atropos segmentation)
if hasattr(img, 'get_data') and hasattr(brain_mask,
'get_data') and not use_fcm:
img = nibabel_to_ants(img)
brain_mask = nibabel_to_ants(brain_mask)
if not use_fcm:
res = img.kmeans_segmentation(3, kmask=brain_mask, mrf=mrf)
avg_intensity = [
np.mean(img.numpy()[prob_img.numpy() > 0.5])
for prob_img in res['probabilityimages']
]
csf_arg = np.argmin(avg_intensity) if contrast.lower() in (
't1', 'flair') else np.argmax(avg_intensity)
csf = res['probabilityimages'][csf_arg].numpy()
if not return_prob:
csf = (csf > csf_thresh).astype(np.float32)
else:
if hasattr(img, 'numpy') and hasattr(brain_mask, 'numpy'):
img = to_nibabel(img)
brain_mask = to_nibabel(brain_mask)
seg = mask.fcm_class_mask(img, brain_mask, hard_seg=True)
avg_intensity = [
np.mean(img.get_fdata()[seg == i]) for i in range(1, 4)
]
csf_arg = np.argmin(avg_intensity) if contrast.lower() in (
't1', 'flair') else np.argmax(avg_intensity)
csf = (seg == (csf_arg + 1)).astype(np.float32)
return csf
def main(args=None):
args = arg_parser().parse_args(args)
if args.verbosity == 1:
level = logging.getLevelName('INFO')
elif args.verbosity >= 2:
level = logging.getLevelName('DEBUG')
else:
level = logging.getLevelName('WARNING')
logging.basicConfig(
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
level=level)
logger = logging.getLogger(__name__)
try:
if not os.path.isdir(args.img_dir):
raise ValueError(
'(-i / --img-dir) argument needs to be a directory of NIfTI images.'
)
if args.mask_dir is not None:
if not os.path.isdir(args.mask_dir):
raise ValueError(
'(-m / --mask-dir) argument needs to be a directory of NIfTI images.'
)
img_fns = io.glob_nii(args.img_dir)
if args.mask_dir is not None:
mask_fns = io.glob_nii(args.mask_dir)
else:
mask_fns = [None] * len(img_fns)
if not os.path.exists(args.output_dir):
logger.info('Making Output Directory: {}'.format(args.output_dir))
os.mkdir(args.output_dir)
hard_seg = not args.memberships
for i, (img_fn, mask_fn) in enumerate(zip(img_fns, mask_fns), 1):
_, base, _ = io.split_filename(img_fn)
logger.info('Creating Mask for Image: {}, ({:d}/{:d})'.format(
base, i, len(img_fns)))
img = io.open_nii(img_fn)
mask = io.open_nii(mask_fn)
tm = fcm_class_mask(img, mask,
hard_seg) if not args.gmm else gmm_class_mask(
img, mask, 't1', False, hard_seg)
tissue_mask = os.path.join(args.output_dir, base + '_tm')
if args.memberships:
classes = ('csf', 'gm', 'wm')
for j, c in enumerate(classes):
io.save_nii(img, tissue_mask + '_' + c + '.nii.gz', tm[...,
j])
else:
io.save_nii(img, tissue_mask + '.nii.gz', tm)
return 0
except Exception as e:
logger.exception(e)
return 1
def find_wm_mask(img, brain_mask, threshold=0.8):
"""
find WM mask using FCM with a membership threshold
Args:
img (nibabel.nifti1.Nifti1Image): target img
brain_mask (nibabel.nifti1.Nifti1Image): brain mask for img
threshold (float): membership threshold
Returns:
wm_mask (nibabel.nifti1.Nifti1Image): white matter mask for img
"""
t1_mem = mask.fcm_class_mask(img, brain_mask)
wm_mask = t1_mem[..., 2] > threshold
wm_mask_nifti = nib.Nifti1Image(wm_mask, img.affine, img.header)
return wm_mask_nifti
def find_tissue_mask(img, brain_mask, threshold=0.8, tissue_type='wm'):
"""
find tissue mask using FCM with a membership threshold
Args:
img (nibabel.nifti1.Nifti1Image): target img
brain_mask (nibabel.nifti1.Nifti1Image): brain mask for img
threshold (float): membership threshold
tissue_type (str): find the mask of this tissue type (wm, gm, or csf)
Returns:
tissue_mask_nifti (nibabel.nifti1.Nifti1Image): tissue mask for img
"""
tissue_to_int = {'csf': 0, 'gm': 1, 'wm': 2}
t1_mem = mask.fcm_class_mask(img, brain_mask)
tissue_mask = t1_mem[..., tissue_to_int[tissue_type]] > threshold
tissue_mask_nifti = nib.Nifti1Image(tissue_mask, img.affine, img.header)
return tissue_mask_nifti
def test_fcm_mask(self):
m = mask.fcm_class_mask(self.img, self.brain_mask, hard_seg=True)
self.assertEqual(len(np.unique(m)), 4)