How to use the pynets.registration.reg_utils function in pynets

def test_align():
    """
    Test align functionality
    """
    # Linear registrattion
    base_dir = str(Path(__file__).parent/"examples")
    anat_dir = base_dir + '/003/anat'
    inp = anat_dir + '/sub-003_T1w_brain.nii.gz'
    ref = anat_dir + '/MNI152_T1_2mm_brain.nii.gz'
    out = anat_dir + '/highres2standard.nii.gz'
    xfm_out = anat_dir + '/highres2standard.mat'

    reg_utils.align(inp, ref, xfm=xfm_out, out=out, dof=12, searchrad=True, bins=256, interp=None, cost="mutualinfo",
                    sch=None, wmseg=None, init=None)
    
    highres2standard_linear = nib.load(out)
    assert highres2standard_linear is not None

# This test has a bak folder in its test_dir. 
    # To replicate test rm data in test_dir and cp from bak
    base_dir = str(Path(__file__).parent/"examples")
    test_dir = base_dir + '/003/test_out/test_check_orient_and_dims'

    # Antomical: 1x1x1mm
    anat_LAS = test_dir + '/anat_LAS/sub-003_T1w_LAS.nii.gz'
    anat_RAS = test_dir + '/anat_RAS/sub-003_T1w_RAS.nii.gz'
    # Diffusion: 2x2x2mm
    dmri_LAS = test_dir + '/dmri_LAS/iso_eddy_corrected_data_denoised_LAS.nii.gz'
    dmri_RAS = test_dir + '/dmri_RAS/iso_eddy_corrected_data_denoised_RAS.nii.gz'
    bvecs_LAS = test_dir + '/dmri_LAS/bvec.orig.bvec'
    bvecs_RAS = test_dir + '/dmri_RAS/bvec.trans.bvec'

    anat_LAStoRAS = reg_utils.check_orient_and_dims(anat_LAS, '2mm', bvecs=None)
    anat_RAStoRAS = reg_utils.check_orient_and_dims(anat_RAS, '2mm', bvecs=None)
    dmri_LAStoRAS, bvecs_LAStoRAS = reg_utils.check_orient_and_dims(dmri_LAS, '1mm', bvecs=bvecs_LAS)
    dmri_RAStoRAS, bvecs_RAStoRAS = reg_utils.check_orient_and_dims(dmri_RAS, '1mm', bvecs=bvecs_RAS)

    anat_LAStoRAS = nib.load(anat_LAStoRAS)

    anat_RAStoRAS = nib.load(anat_RAStoRAS)

    dmri_LAStoRAS = nib.load(dmri_LAStoRAS)

    dmri_RAStoRAS = nib.load(dmri_RAStoRAS)
    
    # Assert that output arrays are identical.
    anat_check = np.allclose(anat_LAStoRAS.affine.astype('int'), anat_RAStoRAS.affine.astype('int'))
    dmri_check = np.allclose(dmri_LAStoRAS.affine.astype('int'), dmri_RAStoRAS.affine.astype('int'))

    # Assert that voxel dimensions in ouputs are correct.

#                                                                                         None else ''),
    #                                                                '%s' % (op.basename(roi).split('.')[0] + '_' if
    #                                                                        roi is not None else ''),
    #                                                                conn_model, '_', target_samples,
    #                                                                '%s' % ("%s%s" %
    #                                                                        ('_' + str(node_size),
    #                                                                         'mm_') if ((node_size != 'parc') and
    #                                                                                    (node_size is not None)) else
    #                                                                        '_'),
    #                                                                'curv', str(curv_thr_list).replace(', ', '_'),
    #                                                                'step', str(step_list).replace(', ', '_'), 'tt-',
    #                                                                track_type,  '_dg-', directget, '_ml-', max_length,
    #                                                                '.png')

    # SyN FA->Template
    [mapping, affine_map, warped_fa] = regutils.wm_syn(template_path, fa_path, dsn_dir)

    tractogram = load_tractogram(streams, fa_img, to_space=Space.RASMM, shifted_origin=True, bbox_valid_check=False)
    fa_img.uncache()
    streamlines = tractogram.streamlines
    warped_fa_img = nib.load(warped_fa)
    warped_fa_affine = warped_fa_img.affine
    warped_fa_shape = warped_fa_img.shape

    streams_in_curr_grid = transform_streamlines(streamlines, warped_fa_affine)

    ref_grid_aff = vox_size*np.eye(4)
    ref_grid_aff[3][3] = 1

    # Create isocenter mapping where we anchor the origin transformation affine
    # to the corner of the FOV by scaling x, y, z offsets according to a multiplicative
    # van der Corput sequence with a base value equal to the voxel resolution

print('Running non-linear registration: T1w-->MNI ...')
                # Use FNIRT to nonlinearly align T1 to MNI template
                regutils.align_nonlinear(self.t1w_brain, self.input_mni, xfm=self.t12mni_xfm_init,
                                         out=self.t1_aligned_mni, warp=self.warp_t1w2mni, ref_mask=self.input_mni_mask)

                # Get warp from MNI -> T1
                regutils.inverse_warp(self.t1w_brain, self.mni2t1w_warp, self.warp_t1w2mni)

                # Get mat from MNI -> T1
                os.system("convert_xfm -omat {} -inverse {}".format(self.mni2t1_xfm_init, self.t12mni_xfm_init))

            except RuntimeError('Error: FNIRT failed!'):
                pass
        else:
            # Falling back to linear registration
            regutils.align(self.t1w_brain, self.input_mni_brain, xfm=self.t12mni_xfm, init=self.t12mni_xfm_init,
                           bins=None, dof=12, cost='mutualinfo', searchrad=True, interp="spline",
                           out=self.t1_aligned_mni, sch=None)
            # Get mat from MNI -> T1
            os.system("convert_xfm -omat {} -inverse {}".format(self.t12mni_xfm, self.mni2t1_xfm))

        # Align T1w-->DWI
        regutils.align(self.fa_path, self.t1w_brain, xfm=self.t1w2dwi_xfm, bins=None, interp="spline", dof=6,
                       cost='mutualinfo', out=None, searchrad=True, sch=None)
        os.system("convert_xfm -omat {} -inverse {}".format(self.dwi2t1w_xfm, self.t1w2dwi_xfm))

        if self.simple is False:
            # Flirt bbr
            try:
                print('Running FLIRT BBR registration: T1w-->DWI ...')
                regutils.align(self.fa_path, self.t1w_brain, wmseg=self.wm_edge, xfm=self.dwi2t1w_bbr_xfm,
                               init=self.dwi2t1w_xfm, bins=256, dof=7, searchrad=True, interp="spline", out=None,

regutils.apply_warp(self.t1w_brain, self.vent_mask_mni, self.vent_mask_t1w, warp=self.mni2t1w_warp,
                                interp='nn', sup=True)

            regutils.apply_warp(self.t1w_brain, self.corpuscallosum, self.corpuscallosum_mask_t1w,
                                warp=self.mni2t1w_warp, interp="nn", sup=True)

        else:
            regutils.applyxfm(self.vent_mask_mni, self.t1w_brain, self.mni2t1_xfm, self.vent_mask_t1w)
            regutils.applyxfm(self.corpuscallosum, self.t1w_brain, self.mni2t1_xfm, self.corpuscallosum_mask_t1w)

        # Applyxfm tissue maps to dwi space
        regutils.applyxfm(self.fa_path, self.vent_mask_t1w, self.t1wtissue2dwi_xfm, self.vent_mask_dwi)
        regutils.applyxfm(self.fa_path, self.csf_mask, self.t1wtissue2dwi_xfm, self.csf_mask_dwi)
        regutils.applyxfm(self.fa_path, self.gm_mask, self.t1wtissue2dwi_xfm, self.gm_in_dwi)
        regutils.applyxfm(self.fa_path, self.wm_mask, self.t1wtissue2dwi_xfm, self.wm_in_dwi)
        regutils.applyxfm(self.fa_path, self.corpuscallosum_mask_t1w, self.t1wtissue2dwi_xfm, self.corpuscallosum_dwi)

        # Threshold WM to binary in dwi space
        thr_img = nib.load(self.wm_in_dwi)
        thr_img = math_img('img > 0.1', img=thr_img)
        nib.save(thr_img, self.wm_in_dwi_bin)

        # Threshold GM to binary in dwi space
        thr_img = nib.load(self.gm_in_dwi)
        thr_img = math_img('img > 0.2', img=thr_img)
        nib.save(thr_img, self.gm_in_dwi_bin)

        # Threshold CSF to binary in dwi space
        thr_img = nib.load(self.csf_mask_dwi)
        thr_img = math_img('img > 0.95', img=thr_img)
        nib.save(thr_img, self.csf_mask_dwi)

"""
        A function to perform alignment from T1w --> MNI and T1w_MNI --> DWI. Uses a local optimisation
        cost function to get the two images close, and then uses bbr to obtain a good alignment of brain boundaries.
        Assumes input dwi is already preprocessed and brain extracted.
        """

        # Create linear transform/ initializer T1w-->MNI
        regutils.align(self.t1w_brain, self.input_mni_brain, xfm=self.t12mni_xfm_init, bins=None, interp="spline",
                       out=None, dof=12, cost='mutualinfo', searchrad=True)

        # Attempt non-linear registration of T1 to MNI template
        if self.simple is False:
            try:
                print('Running non-linear registration: T1w-->MNI ...')
                # Use FNIRT to nonlinearly align T1 to MNI template
                regutils.align_nonlinear(self.t1w_brain, self.input_mni, xfm=self.t12mni_xfm_init,
                                         out=self.t1_aligned_mni, warp=self.warp_t1w2mni, ref_mask=self.input_mni_mask)

                # Get warp from MNI -> T1
                regutils.inverse_warp(self.t1w_brain, self.mni2t1w_warp, self.warp_t1w2mni)

                # Get mat from MNI -> T1
                os.system("convert_xfm -omat {} -inverse {}".format(self.mni2t1_xfm_init, self.t12mni_xfm_init))

            except RuntimeError('Error: FNIRT failed!'):
                pass
        else:
            # Falling back to linear registration
            regutils.align(self.t1w_brain, self.input_mni_brain, xfm=self.t12mni_xfm, init=self.t12mni_xfm_init,
                           bins=None, dof=12, cost='mutualinfo', searchrad=True, interp="spline",
                           out=self.t1_aligned_mni, sch=None)
            # Get mat from MNI -> T1

uatlas_filled = "%s%s%s%s" % (self.anat_path, '/', atlas, "_filled.nii.gz")
        os.system("fslmaths {} -add {} -mas {} {}".format(self.input_mni_brain, uatlas, self.input_mni_mask,
                                                          uatlas_filled))

        regutils.align(uatlas_filled, self.t1_aligned_mni, init=None, xfm=self.atlas2t1wmni_xfm_init,
                       out=None, dof=12, searchrad=True, interp="nearestneighbour", cost='mutualinfo')

        if uatlas_parcels is not None:
            regutils.applyxfm(self.t1_aligned_mni, uatlas_parcels, self.atlas2t1wmni_xfm_init, aligned_atlas_t1mni,
                              interp="nearestneighbour")
        else:
            regutils.applyxfm(self.t1_aligned_mni, uatlas, self.atlas2t1wmni_xfm_init, aligned_atlas_t1mni,
                              interp="nearestneighbour")

        try:
            regutils.apply_warp(self.t1_aligned_mni, self.gm_mask_thr, gm_mask_mni, warp=self.warp_t1w2mni,
                                xfm=self.t12mni_xfm_init, interp='nn', sup=True)
        except:
            regutils.applyxfm(self.t1_aligned_mni, self.gm_mask_thr, self.t12mni_xfm_init, gm_mask_mni,
                              interp="nearestneighbour")

        # Set intensities to int
        atlas_img = nib.load(aligned_atlas_t1mni)
        gm_mask_img_res = resample_img(nib.load(gm_mask_mni), target_affine=atlas_img.affine,
                                       target_shape=atlas_img.shape)
        nib.save(gm_mask_img_res, gm_mask_mni_atlas_res)
        os.system("fslmaths {} -bin {}".format(gm_mask_mni_atlas_res, gm_mask_mni_atlas_res))
        os.system("fslmaths {} -mas {} {}".format(aligned_atlas_t1mni, gm_mask_mni_atlas_res, aligned_atlas_t1mni_gm))

        gm_mask_img_res.uncache()
        atlas_img.uncache()

"""

        # Create linear transform/ initializer T1w-->MNI
        regutils.align(self.t1w_brain, self.input_mni_brain, xfm=self.t12mni_xfm_init, bins=None, interp="spline",
                       out=None, dof=12, cost='mutualinfo', searchrad=True)

        # Attempt non-linear registration of T1 to MNI template
        if self.simple is False:
            try:
                print('Running non-linear registration: T1w-->MNI ...')
                # Use FNIRT to nonlinearly align T1 to MNI template
                regutils.align_nonlinear(self.t1w_brain, self.input_mni, xfm=self.t12mni_xfm_init,
                                         out=self.t1_aligned_mni, warp=self.warp_t1w2mni, ref_mask=self.input_mni_mask)

                # Get warp from MNI -> T1
                regutils.inverse_warp(self.t1w_brain, self.mni2t1w_warp, self.warp_t1w2mni)

                # Get mat from MNI -> T1
                os.system("convert_xfm -omat {} -inverse {}".format(self.mni2t1_xfm_init, self.t12mni_xfm_init))

            except RuntimeError('Error: FNIRT failed!'):
                pass
        else:
            # Falling back to linear registration
            regutils.align(self.t1w_brain, self.input_mni_brain, xfm=self.t12mni_xfm, init=self.t12mni_xfm_init,
                           bins=None, dof=12, cost='mutualinfo', searchrad=True, interp="spline",
                           out=self.t1_aligned_mni, sch=None)
            # Get mat from MNI -> T1
            os.system("convert_xfm -omat {} -inverse {}".format(self.t12mni_xfm, self.mni2t1_xfm))

        # Align T1w-->DWI
        regutils.align(self.fa_path, self.t1w_brain, xfm=self.t1w2dwi_xfm, bins=None, interp="spline", dof=6,

aligned_atlas_skull_parcels = None

                # Map atlas in t1w space to dwi space
                if uatlas_parcels is not None:
                    regutils.applyxfm(self.fa_path, aligned_atlas_skull_parcels, self.t1wtissue2dwi_xfm,
                                      dwi_aligned_atlas, interp="nearestneighbour")
                else:
                    regutils.applyxfm(self.fa_path, aligned_atlas_skull, self.t1wtissue2dwi_xfm,
                                      dwi_aligned_atlas, interp="nearestneighbour")
            except:
                print("Warning: Atlas is not in correct dimensions, or input is low quality,\nusing linear template "
                      "registration.")

                # Combine our linear transform from t1w to template with our transform from dwi to t1w space to get a
                # transform from atlas ->(-> t1w ->)-> dwi
                regutils.combine_xfms(self.xfm_atlas2t1w_init, self.t1wtissue2dwi_xfm, self.temp2dwi_xfm)

                if uatlas_parcels is not None:
                    aligned_atlas_t1mni_parcels = "%s%s%s%s" % (self.anat_path, '/', atlas, "_t1w_mni_parcels.nii.gz")
                    regutils.applyxfm(self.t1_aligned_mni, uatlas_parcels, self.atlas2t1mni_xfm_init,
                                      aligned_atlas_t1mni_parcels, interp="nearestneighbour")
                    # Apply linear transformation from template to dwi space
                    regutils.applyxfm(self.fa_path, aligned_atlas_t1mni_parcels, self.temp2dwi_xfm, dwi_aligned_atlas,
                                      interp="nearestneighbour")
                    aligned_atlas_t1mni = aligned_atlas_t1mni_parcels
                else:
                    regutils.applyxfm(self.t1_aligned_mni, uatlas, self.atlas2t1mni_xfm_init,
                                      aligned_atlas_t1mni, interp="nearestneighbour")
                    # Apply linear transformation from template to dwi space
                    regutils.applyxfm(self.fa_path, aligned_atlas_t1mni, self.temp2dwi_xfm, dwi_aligned_atlas,
                                      interp="nearestneighbour")
        else:

def gen_tissue(self):
        """
        A function to segment and threshold tissue types from T1w.
        """
        # Segment the t1w brain into probability maps
        maps = regutils.segment_t1w(self.t1w_brain, self.map_path)
        self.wm_mask = maps['wm_prob']
        self.gm_mask = maps['gm_prob']
        self.csf_mask = maps['csf_prob']

        # Threshold WM to binary in dwi space
        t_img = nib.load(self.wm_mask)
        mask = math_img('img > 0.2', img=t_img)
        mask.to_filename(self.wm_mask_thr)

        # Threshold T1w brain to binary in anat space
        t_img = nib.load(self.t1w_brain)
        mask = math_img('img > 0.0', img=t_img)
        mask.to_filename(self.t1w_brain_mask)

        # Extract wm edge
        os.system("fslmaths {} -edge -bin -mas {} {}".format(self.wm_mask_thr, self.wm_mask_thr, self.wm_edge))