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)