write_with_smooth_deface.py

def write_anat(bids_root, subject, t1w, session=None, acquisition=None,
               raw=None, trans=None, deface=False, overwrite=False,
               verbose=False):
    """Put anatomical MRI data into a BIDS format.

    Given a BIDS directory and a T1 weighted MRI scan for a certain subject,
    format the MRI scan to be in BIDS format and put it into the correct
    location in the bids_dir. If a transformation matrix is supplied, a
    sidecar JSON file will be written for the T1 weighted data.

    Parameters
    ----------
    bids_root : str
        Path to root of the BIDS folder
    subject : str
        Subject label as in 'sub-<label>', for example: '01'
    t1w : str | nibabel image object
        Path to a T1 weighted MRI scan of the subject. Can be in any format
        readable by nibabel. Can also be a nibabel image object of a T1
        weighted MRI scan. Will be written as a .nii.gz file.
    session : str | None
        The session for `t1w`. Corresponds to "ses"
    acquisition: str | None
        The acquisition parameters for `t1w`. Corresponds to "acq"
    raw : instance of Raw | None
        The raw data of `subject` corresponding to `t1w`. If `raw` is None,
        `trans` has to be None as well
    trans : instance of mne.transforms.Transform | str | None
        The transformation matrix from head coordinates to MRI coordinates. Can
        also be a string pointing to a .trans file containing the
        transformation matrix. If None, no sidecar JSON file will be written
        for `t1w`
    deface : bool | tuple (float, float, str, bool) | None
        If True, deface with default parameters, if parameters are provided,
        the order is inset, theta, method, plot. Here `inset`
        is how far back as a fraction of mri space to start defacing
        relative to the nasion (default 0.2), `theta` is the angle of
        the defacing shear (default 35 degrees), `method` accepts
        `smooth` to apply a guassian kernal to the masked
        face area and `erase` to set the image values in the face area
        to zero (default `smooth`, `plot_on` is whether or not to
        plot the results (default False). If None, no defacing
        will be performed.
    overwrite : bool
        Whether to overwrite existing files or data in files.
        Defaults to False.
        If overwrite is True, any existing files with the same BIDS parameters
        will be overwritten with the exception of the `participants.tsv` and
        `scans.tsv` files. For these files, parts of pre-existing data that
        match the current data will be replaced.
        If overwrite is False, no existing data will be overwritten or
        replaced.
    verbose : bool
        If verbose is True, this will print a snippet of the sidecar files. If
        False, no content will be printed.

    Returns
    -------
    anat_dir : str
        Path to the anatomical scan in the `bids_dir`

    """
    if not has_nibabel():  # pragma: no cover
        raise ImportError('This function requires nibabel.')
    import nibabel as nib

    if deface and trans is None:
        raise ValueError('The raw object and trans must be provided to '
                         'deface the T1')

    if isinstance(deface, list) or isinstance(deface, tuple):
        inset, theta, method, plot_on = deface
        assert isinstance(inset, float) or isinstance(inset, int)
        assert isinstance(theta, float) or isinstance(inset, int)
        assert isinstance(method, str)
        assert isinstance(plot_on, bool)
    elif deface:
        inset, theta, method, plot_on = (0.2, 35., 'smooth', False)

    # Make directory for anatomical data
    anat_dir = op.join(bids_root, 'sub-{}'.format(subject))
    # Session is optional
    if session is not None:
        anat_dir = op.join(anat_dir, 'ses-{}'.format(session))
    anat_dir = op.join(anat_dir, 'anat')
    if not op.exists(anat_dir):
        os.makedirs(anat_dir)

    # Try to read our T1 file and convert to MGH representation
    if isinstance(t1w, str):
        t1w = nib.load(t1w)
    elif type(t1w) not in nib.all_image_classes:
        raise ValueError('`t1w` must be a path to a T1 weighted MRI data file '
                         ', or a nibabel image object, but it is of type '
                         '"{}"'.format(type(t1w)))

    t1w = nib.Nifti1Image(t1w.dataobj, t1w.affine)
    # XYZT_UNITS = NIFT_UNITS_MM (10 in binary or 2 in decimal)
    # seems to be the default for Nifti files
    # https://nifti.nimh.nih.gov/nifti-1/documentation/nifti1fields/nifti1fields_pages/xyzt_units.html
    if t1w.header['xyzt_units'] == 0:
        t1w.header['xyzt_units'] = np.array(10, dtype='uint8')

    # Now give the NIfTI file a BIDS name and write it to the BIDS location
    t1w_basename = make_bids_basename(subject=subject, session=session,
                                      acquisition=acquisition, prefix=anat_dir,
                                      suffix='T1w.nii.gz')

    # Check if we have necessary conditions for writing a sidecar JSON
    if trans is not None:

        # get trans and ensure it is from head to MRI
        trans, _ = _get_trans(trans, fro='head', to='mri')

        if not isinstance(raw, BaseRaw):
            raise ValueError('`raw` must be specified if `trans` is not None')

        # Prepare to write the sidecar JSON
        # extract MEG landmarks
        coords_dict = _extract_landmarks(raw.info['dig'])
        meg_landmarks = np.asarray((coords_dict['LPA'],
                                    coords_dict['NAS'],
                                    coords_dict['RPA']))

        # Transform MEG landmarks into MRI space, adjust units by * 1e3
        mri_landmarks = apply_trans(trans, meg_landmarks, move=True) * 1e3

        # Get landmarks in voxel space, using the mgh version of our T1 data
        t1_mgh = nib.MGHImage(t1w.dataobj, t1w.affine)
        vox2ras_tkr = t1_mgh.header.get_vox2ras_tkr()
        ras2vox_tkr = np.linalg.inv(vox2ras_tkr)
        mri_landmarks = apply_trans(ras2vox_tkr, mri_landmarks)  # in vox

        # Write sidecar.json
        t1w_json = dict()
        t1w_json['AnatomicalLandmarkCoordinates'] = \
            {'LPA': list(mri_landmarks[0, :]),
             'NAS': list(mri_landmarks[1, :]),
             'RPA': list(mri_landmarks[2, :])}
        fname = t1w_basename.replace('.nii.gz', '.json')
        if op.isfile(fname) and not overwrite:
            raise IOError('Wanted to write a file but it already exists and '
                          '`overwrite` is set to False. File: "{}"'
                          .format(fname))
        _write_json(fname, t1w_json, overwrite, verbose)

        if deface:
            # x: L/R L+, y: S/I I+, z: A/P A+
            t1w_data = t1w.get_data().copy()
            indices = np.meshgrid(np.arange(t1w_data.shape[0]),
                                  np.arange(t1w_data.shape[1]),
                                  np.arange(t1w_data.shape[2]),
                                  indexing='ij')
            indices = np.array(indices)  # e.g., (3, 86, 86, 86)
            indices = np.transpose(indices, [1, 2, 3, 0])  # (86, 86, 86, 3)
            indices = indices.reshape(-1, 3)  # (86 * 86 * 86, 3)

            mri_landmarks2 = apply_trans(t1w.affine, mri_landmarks)
            meg_trans = get_ras_to_neuromag_trans(*mri_landmarks2[[1, 0, 2]])

            indices = apply_trans(t1w.affine, indices)
            indices = apply_trans(meg_trans, indices)
            trans_y = -mri_landmarks2[1, 1] + t1w_data.shape[2] * inset
            indices = apply_trans(translation(y=trans_y), indices)
            indices = apply_trans(rotation(x=-np.deg2rad(theta)), indices)
            coords = indices.reshape(t1w.shape + (3,))
            mask = (coords[..., 2] < 0)

            if method == 'smooth':
                # https://docs.scipy.org/doc/scipy-0.16.1/reference/generated/scipy.ndimage.filters.gaussian_filter.html
                from scipy.ndimage import gaussian_filter
                for layer in range(5):
                    t1w_data[mask] = gaussian_filter(t1w_data,
                                                     sigma=0.5 * layer,
                                                     order=0)[mask]
                    z = t1w_data.shape[2] * 0.1
                    indices = apply_trans(translation(z=z),
                                          indices)
                    coords = indices.reshape(t1w.shape + (3,))
                    mask = (coords[..., 2] < 0)

            elif method == 'erase':
                t1w_data[mask] = 0.
            else:
                raise ValueError('Deface argument %s not recognized' % method)

            t1w = nib.Nifti1Image(t1w_data, t1w.affine, t1w.header)
            if plot_on:
                fig = t1w.orthoview()
                fig.show()

    # Save anatomical data
    if op.exists(t1w_basename):
        if overwrite:
            os.remove(t1w_basename)
        else:
            raise IOError('Wanted to write a file but it already exists and '
                          '`overwrite` is set to False. File: "{}"'
                          .format(t1w_basename))

    nib.save(t1w, t1w_basename)

    return anat_dir