Plot atlas on brain surface

plot_atlas_on_surf.py

from typing import Literal, List

import surfplot
import matplotlib as plt
import nibabel as nib
import numpy as np
import pandas as pd


def plot_atlas_on_surf(atlas_path: str,
                       data: pd.DataFrame,
                       cmap: str = 'viridis',
                       surf_type: Literal['white', 'pial', 'inflated', 'sphere', 'medial', 'sulc', 'vaavg'] = 'white',
                       views: List = ['lateral', 'medial', 'ventral']):
    import neuromaps.datasets
    import neuromaps.transforms
    import neuromaps.images
    
    # path to surface mesh
    surfaces = neuromaps.datasets.fetch_fsaverage(density='41k')
    lh, rh = surfaces[surf_type]
    
    # path to atlas
    volumetric_atlas = nib.load(atlas_path)
    atlas_label = volumetric_atlas.get_fdata()
    
    atlas_data = np.zeros(atlas_label.shape)
    for _, row in data.iterrows():
        atlas_data[atlas_label == row['idx']] = row['value']
    
    new_atlas_img = nib.Nifti1Image(atlas_data, affine=volumetric_atlas.affine, header=volumetric_atlas.header)
    
    # transform volumetric atlas to surface
    plot_lh_data, plot_rh_data = neuromaps.transforms.mni152_to_fsaverage(
        neuromaps.images.load_nifti(new_atlas_img),
        fsavg_density='41k',
        method='nearest',
    )
    
    p = surfplot.Plot(
        surf_lh=lh,
        surf_rh=rh,
        views=views,
        layout='grid',
        zoom=1.2,
        size=np.array([1000]) * (4, 6),
    )
    p.add_layer({'left': plot_lh_data, 'right': plot_rh_data}, cmap=cmap)
    fig = p.build(figsize=(10, 10))
    
    return fig