Inflate a brain from the pial surface to visualize sEEG contacts

inflate_brain.py

import os
import os.path as op
import numpy as np
import mne
import imageio

misc_path = mne.datasets.misc.data_path()
sample_path = mne.datasets.sample.data_path()
subjects_dir = misc_path / 'seeg'
subject = 'sample_seeg'

raw = mne.io.read_raw(misc_path / 'seeg' / 'sample_seeg_ieeg.fif')
trans = mne.coreg.estimate_head_mri_t('sample_seeg', subjects_dir)

view_kwargs = dict(azimuth=120, elevation=100, distance=600,
                   focalpoint=(0, 0, -15))

proj_info = mne.preprocessing.ieeg.project_sensors_onto_inflated(
    raw.info, trans, subject, subjects_dir=subjects_dir)

surf_data = dict(lh=dict(), rh=dict())
x_dir = np.array([1., 0., 0.])
for hemi in ('lh', 'rh'):
    for surf in ('pial', 'inflated', 'curv'):
        for img in ('', '.T1', '.T2', ''):
            surf_fname = op.join(subjects_dir, subject, 'surf',
                                 f'{hemi}.{surf}')
            if op.isfile(surf_fname):
                break
        if surf == 'curv':
            surf_data[hemi]['curv'] = np.array(mne.surface.read_curvature(
                surf_fname, binary=False) > 0, np.int64)
        else:
            coords, faces = mne.surface.read_surface(surf_fname)
            if surf == 'inflated':
                x_ = coords @ x_dir
                coords -= np.max(x_) * x_dir if hemi == 'lh' else \
                    np.min(x_) * x_dir
            surface = dict(rr=coords, tris=faces)
            mne.surface.complete_surface_info(
                surface, copy=False, verbose=False, do_neighbor_tri=False)
            surf_data[hemi][surf] = surface['rr'], surface['tris'], surface['nn']


for hemi in ('lh', 'rh'):
    surf_data[hemi]['vectors'] = \
        surf_data[hemi]['inflated'][0] - surf_data[hemi]['pial'][0]
    surf_data[hemi]['normal_vectors'] = \
        surf_data[hemi]['inflated'][2] - surf_data[hemi]['pial'][2]

montage_start = raw.get_montage()
montage_start.apply_trans(trans)
ch_pos_start = montage_start.get_positions()['ch_pos']
montage_end = proj_info.get_montage()
montage_end.apply_trans(trans)
ch_pos_end = montage_end.get_positions()['ch_pos']

ch_pos_vectors = dict()
for ch in raw.ch_names:
    ch_pos_vectors[ch] = ch_pos_end[ch] - ch_pos_start[ch]


brain = mne.viz.Brain(subject, subjects_dir=subjects_dir,
                      cortex='low_contrast', alpha=0.5, background='white')
# brain.add_sensors(raw.info, trans=trans)
brain.show_view(**view_kwargs)

brain.add_annotation('aparc.a2009s', borders=False, alpha=0.5)

sensor_meshes = dict()
sensor_locs = dict()
for ch in raw.ch_names:
    loc = ch_pos_start[ch].copy()
    sensor_mesh = brain._renderer.sphere(loc * 1000, 'yellow', 3)[1]
    sensor_meshes[ch] = sensor_mesh
    sensor_locs[ch] = sensor_mesh.points.copy()

images = [brain.screenshot()] * 5

n_steps = 101
for t in np.linspace(0, 1, n_steps):

    for hemi in ('lh', 'rh'):
        coords, faces, nn = surf_data[hemi]['pial']
        coords = coords.copy()
        coords += surf_data[hemi]['vectors'] * t
        nn = nn.copy()
        nn += surf_data[hemi]['normal_vectors'] * t
        brain._renderer.plotter.update_coordinates(
            coords, brain._layered_meshes[hemi]._polydata, render=False)
        brain._layered_meshes[hemi]._polydata.point_data.active_normals = nn
        brain._layered_meshes[hemi].update_overlay('curv', opacity=0.5 + t / 2)

    for ch, sensor_mesh in sensor_meshes.items():
        loc = sensor_locs[ch].copy()
        loc += ch_pos_vectors[ch] * t * 1000
        brain._renderer.plotter.update_coordinates(
            loc, sensor_mesh, render=False)

    brain._renderer._update()
    images.append(brain.screenshot())

for i in range(5):
    images.append(images[-1])

brain.close()
imageio.mimwrite('inflated.mp4', images, fps=24)