def test_gmm_mask(self):
wm_peak = mask.gmm_class_mask(self.img,
self.brain_mask,
return_wm_peak=True)
self.assertAlmostEqual(wm_peak, 673952.0553833192)
m = mask.gmm_class_mask(self.img,
self.brain_mask,
return_wm_peak=False,
hard_seg=True)
self.assertEqual(len(np.unique(m)), 4)
m = mask.gmm_class_mask(self.img,
self.brain_mask,
return_wm_peak=False,
hard_seg=False)
self.assertEqual(m.shape[3], 3)
def process(image, brain_mask, args, logger):
img = io.open_nii(image)
mask = io.open_nii(brain_mask)
dirname, base, ext = io.split_filename(image)
if args.output_dir is not None:
dirname = args.output_dir
if not os.path.exists(dirname):
logger.info('Making output directory: {}'.format(dirname))
os.mkdir(dirname)
if args.find_background_mask:
bg_mask = background_mask(img)
bgfile = os.path.join(dirname, base + '_bgmask' + ext)
io.save_nii(bg_mask, bgfile, is_nii=True)
if args.wm_peak is not None:
logger.info('Loading WM peak: ', args.wm_peak)
peak = float(np.load(args.wm_peak))
else:
peak = gmm_class_mask(img, brain_mask=mask, contrast=args.contrast)
if args.save_wm_peak:
np.save(os.path.join(dirname, base + '_wmpeak.npy'), peak)
normalized = gmm.gmm_normalize(img, mask, args.norm_value, args.contrast,
args.background_mask, peak)
outfile = os.path.join(dirname, base + '_gmm' + ext)
logger.info('Normalized image saved: {}'.format(outfile))
io.save_nii(normalized, outfile, is_nii=True)
def gmm_normalize(img, brain_mask=None, norm_value=1, contrast='t1', bg_mask=None, wm_peak=None):
"""
normalize the white matter of an image using a GMM to find the tissue classes
Args:
img (nibabel.nifti1.Nifti1Image): target MR image
brain_mask (nibabel.nifti1.Nifti1Image): brain mask for img
norm_value (float): value at which to place the WM mean
contrast (str): MR contrast type for img
bg_mask (nibabel.nifti1.Nifti1Image): if provided, use to zero bkgd
wm_peak (float): previously calculated WM peak
Returns:
normalized (nibabel.nifti1.Nifti1Image): gmm wm peak normalized image
"""
if wm_peak is None:
wm_peak = gmm_class_mask(img, brain_mask=brain_mask, contrast=contrast)
img_data = img.get_data()
logger.info('Normalizing Data...')
norm_data = (img_data/wm_peak)*norm_value
norm_data[norm_data < 0.1] = 0.0
if bg_mask is not None:
logger.info('Applying background mask...')
masked_image = norm_data * bg_mask.get_data()
else:
masked_image = norm_data
normalized = nib.Nifti1Image(masked_image, img.affine, img.header)
return normalized
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