blender_visualization_subcortex_with_glass_brain.py

# Blender executed via:
# PYTHONPATH=/home/sina/miniconda3/envs/blender_env/bin/python blender  --python-use-system-env

# Notes:
# Setting up blender to use a conda virtualenv:
# https://stackoverflow.com/questions/70639689/how-to-use-the-anaconda-environment-on-blender
# Blender Python API:
# https://docs.blender.org/api/current/info_quickstart.html
# Blender Python tips and tricks:
# https://docs.blender.org/api/current/info_tips_and_tricks.html
# Blender can be built from source to enable direct python execution:
# https://developer.blender.org/docs/handbook/building_blender/python_module/
# Building blender:
# https://developer.blender.org/docs/handbook/building_blender/
# Idea: conda can be used to build blender for a python package:
# https://docs.aws.amazon.com/deadline-cloud/latest/developerguide/create-conda-recipe-blender.html
# https://github.com/aws-deadline/deadline-cloud-samples/blob/mainline/conda_recipes/blender-4.2/recipe/meta.yaml

import bpy
import bmesh
import math
import numpy as np
import nibabel as nib
import matplotlib.pyplot as plt
from Connectome_Spatial_Smoothing import CSS as css
from cerebro import cerebro_brain_utils as cbu
from cerebro import cerebro_brain_viewer as cbv


# function to load brain surface geometries
def load_cortical_meshes(surface):
    left_surface_file, right_surface_file = cbu.get_left_and_right_GIFTI_template_surface(surface)
    left_vertices, left_triangles = cbu.load_GIFTI_surface(left_surface_file)
    right_vertices, right_triangles = cbu.load_GIFTI_surface(right_surface_file)
    return left_vertices, left_triangles, right_vertices, right_triangles

# function to read msa atlas colors
def read_msa_atlas_information(file_path):
    with open(file_path, 'r') as file:
        lines = file.readlines()

    # Strip newlines and concatenate every two lines, then split
    concatenated_lines = [(lines[i].strip() + ' ' + lines[i+1].strip()).split() for i in range(0, len(lines)-1, 2) if '7Networks_' not in lines[i]]

    # Now make dict
    information = {
        float(line[1]): {
            "name": line[0],
            "r": float(line[2])/255,
            "g": float(line[3])/255,
            "b": float(line[4])/255,
            "a": float(line[5])/255}
        for line in concatenated_lines
    }
    
    return information


# function to clean the blender scene
def blender_clean_up():
    # make sure the active object is not in Edit Mode
    if bpy.context.active_object and bpy.context.active_object.mode == "EDIT":
        bpy.ops.object.editmode_toggle()

    # make sure non of the objects are hidden from the viewport, selection, or disabled
    for obj in bpy.data.objects:
        obj.hide_set(False)
        obj.hide_select = False
        obj.hide_viewport = False

    # select all the object and delete them (just like pressing A + X + D in the viewport)
    bpy.ops.object.select_all(action="SELECT")
    bpy.ops.object.delete()

    # find all the collections and remove them
    collection_names = [col.name for col in bpy.data.collections]
    for name in collection_names:
        bpy.data.collections.remove(bpy.data.collections[name])

    # in the case when you modify the world shader
    # delete and recreate the world object
    world_names = [world.name for world in bpy.data.worlds]
    for name in world_names:
        bpy.data.worlds.remove(bpy.data.worlds[name])
    # create a new world data block
    bpy.ops.world.new()
    bpy.context.scene.world = bpy.data.worlds["World"]

    bpy.ops.outliner.orphans_purge(do_local_ids=True, do_linked_ids=True, do_recursive=True)


# function to adjust animation frame setup
def frame_setup(start_frame, end_frame):
    # Set the start and end frames of the animation
    bpy.context.scene.frame_start = start_frame
    bpy.context.scene.frame_end = end_frame


# function to create a mesh
def create_mesh_object(vertices, faces, mesh_name, object_name):
    # Create a new mesh and object
    mesh = bpy.data.meshes.new(mesh_name)
    obj = bpy.data.objects.new(object_name, mesh)

    # Link the object to the current scene collection
    bpy.context.collection.objects.link(obj)
    
    # Set the object as active and select it
    bpy.context.view_layer.objects.active = obj
    obj.select_set(True)
    
    # Create the mesh from the vertices and faces
    mesh.from_pydata(vertices, [], faces)
    mesh.update()  # Update the mesh to make sure it's updated in the scene

    # Return the created object
    return mesh, obj

# function to apply smooth shading to object
def shade_smooth(obj):
    # Ensure the object is selected and active
    bpy.context.view_layer.objects.active = obj
    obj.select_set(True)
    
    # Apply smooth shading
    bpy.ops.object.shade_smooth()


# function to turn data to colors
def generate_vertex_colors(data, surface_model, colormap=plt.cm.coolwarm, alpha=1.0):
    # Normalize the data and map to colors
    scaled_data = (data - data.min()) / (data.max() - data.min())
    vertex_data = np.zeros(surface_model.surface_number_of_vertices) * np.nan
    vertex_data[np.array(surface_model.vertex_indices)] = scaled_data[
        surface_model.index_offset:surface_model.index_offset+surface_model.index_count
    ]
    vertex_colors = colormap(vertex_data)
    # adjust alpha
    vertex_colors[:, 3] = alpha
    return vertex_colors


# Function to update vertex colors on a mesh
def update_vertex_colors(mesh, vertex_colors):
    color_layer = mesh.vertex_colors.active.data
    for loop_index, loop in enumerate(mesh.loops):
        vertex_index = loop.vertex_index
        color_layer[loop_index].color = vertex_colors[vertex_index]


# Function to add vertex colors to a mesh's vertex color layer
def create_layer(mesh, layer_name):
    # Create a new vertex color layer or use an existing one
    if layer_name not in mesh.vertex_colors:
        mesh.vertex_colors.new(name=layer_name)        
    vertex_color_layer = mesh.vertex_colors[layer_name]
    mesh.vertex_colors.active_index = [x[0] for x in mesh.vertex_colors.items()].index(layer_name)
    return vertex_color_layer


# Function to add vertex colors to a mesh's vertex color layer
def add_vertex_colors_to_layer(mesh, vertex_colors, layer_name):
    create_layer(mesh, layer_name)
    update_vertex_colors(mesh, vertex_colors)
    mesh.update()


# Function to create material and assign color to it
def assign_object_material_with_color(obj, color_layer_name, metallic, roughness):
    # Create a new material if it doesn't exist
    material = bpy.data.materials.new("VertexColorMaterial")
    material.use_nodes = True
    nodes = material.node_tree.nodes
    links = material.node_tree.links

    # Clear existing nodes
    for node in nodes:
        nodes.remove(node)

    # Create nodes: Vertex Color node, Principled BSDF, and Material Output
    vertex_color_node = nodes.new(type='ShaderNodeVertexColor')
    vertex_color_node.layer_name = color_layer_name

    principled_bsdf = nodes.new(type='ShaderNodeBsdfPrincipled')
    material_output = nodes.new(type='ShaderNodeOutputMaterial')

    # Connect vertex color layer to the Principled BSDF shader and then to the material output
    links.new(vertex_color_node.outputs['Color'], principled_bsdf.inputs['Base Color'])
    links.new(vertex_color_node.outputs['Alpha'], principled_bsdf.inputs['Alpha'])
    links.new(principled_bsdf.outputs['BSDF'], material_output.inputs['Surface'])

    # Other material properties
    principled_bsdf.inputs["Metallic"].default_value = metallic
    principled_bsdf.inputs["Roughness"].default_value = roughness

    # Assign the material to the object
    if len(obj.data.materials) > 0:
        obj.data.materials[0] = material
    else:
        obj.data.materials.append(material)

    return material


# Function to create object
def create_data_object(name, object_data):
    obj = bpy.data.objects.new(name, object_data)
    bpy.context.collection.objects.link(obj)
    return obj


# Function to create empty
def create_empty(name, location):
    # Create an empty object
    bpy.ops.object.empty_add(type='PLAIN_AXES', location=location)
    empty = bpy.context.object
    empty.name = name
    return empty


# Function to create plane
def create_plane(name, scale, location, rotation, subdivision_kws={}):
    # Create a plane
    bpy.ops.mesh.primitive_plane_add(location=location, rotation=rotation)
    plane = bpy.context.object
    plane.name = name
    plane.scale[0] = scale[0] / 2
    plane.scale[1] = scale[1] / 2

    # subdivide if needed
    if subdivision_kws['subdivide']:
        # Enter edit mode to further subdivide
        bpy.ops.object.mode_set(mode='EDIT')

        bm = bmesh.from_edit_mesh(plane.data)

        # Subdivide the mesh further to get 80 sections
        bmesh.ops.subdivide_edges(bm, edges=bm.edges, cuts=subdivision_kws['cuts'], use_grid_fill=True,)
        bmesh.update_edit_mesh(plane.data)
        bpy.ops.object.mode_set(mode='OBJECT')

    return plane


# Function to create a brain like material
def create_text_object(name, location, rotation, text, align_x="CENTER", align_y="CENTER", font=None):
    bpy.ops.object.text_add(location=location, rotation=rotation)
    text_obj = bpy.context.object
    text_obj.name = name

    # Set the text content
    text_obj.data.body = text
    
    # Align the text within its local coordinates
    text_obj.data.align_x = align_x
    text_obj.data.align_y = align_y

    # Optional: Load a custom font
    if font:
        text_obj.data.font = bpy.data.fonts.load(font)

    return text_obj


# Function to group a subset of plane vertices
def group_top_vertices(plane, group_name, epsilon=0.01):
    # Create a vertex group for the top central vertices
    vertex_group = plane.vertex_groups.new(name=group_name)
    
    # Get the top vertices (considering Z direction)
    top_z_coord = np.min([v.co.x for v in plane.data.vertices])
    max_y_coord = np.max([v.co.y for v in plane.data.vertices])
    mid_y_coord = np.mean([v.co.y for v in plane.data.vertices])
    y_selection_lim = (max_y_coord - mid_y_coord) / 3
    top_verts = [
        v for v in plane.data.vertices
        if (
            (np.abs(v.co.x - top_z_coord) < epsilon) and
            (np.abs(v.co.y - mid_y_coord) <= y_selection_lim)
        )
    ]
    
    # Assign the top vertices to the vertex group
    for v in top_verts:
        vertex_group.add([v.index], 1.0, 'ADD')

    return vertex_group


# Function to set a hook modifier on vertex group
def set_hook_modifier(object_to_be_hooked, vertex_group_name, hook_object):
    # Add a hook modifier to the plane and assign the top vertices
    hook_modifier = object_to_be_hooked.modifiers.new(name="Hook", type='HOOK')
    hook_modifier.object = hook_object
    hook_modifier.vertex_group = vertex_group_name
    return hook_modifier


# Function to add cloth physics to object
def add_cloth_physics(object, pin_vertex_group=None, vertex_mass=1):
    # Ensure the object is a mesh
    if object.type != 'MESH':
        print(f"Error: {object.name} is not a mesh object.")
        return
    
    # Add Cloth modifier
    cloth_modifier = object.modifiers.new(name="Cloth", type='CLOTH')
    
    # Optional: Set the pinning vertex group (if specified)
    if pin_vertex_group is not None and pin_vertex_group in object.vertex_groups:
        cloth_modifier.settings.vertex_group_mass = pin_vertex_group
    
    # Default settings for cloth simulation (you can customize these as needed)
    cloth_modifier.settings.quality = 30  # Quality of the simulation (higher is slower but more accurate)
    cloth_modifier.settings.time_scale = 5  # Speed Multiplier
    cloth_modifier.settings.mass = vertex_mass  # Mass of the cloth
    cloth_modifier.settings.air_damping = 1  # Air resistance
    cloth_modifier.settings.tension_stiffness = 10  # Stiffness of the cloth
    cloth_modifier.settings.compression_stiffness = 10
    cloth_modifier.settings.shear_stiffness = 5
    cloth_modifier.settings.bending_stiffness = 0.2

    # Enable self-collision to prevent parts of the cloth from intersecting itself
    cloth_modifier.collision_settings.use_self_collision = True
    cloth_modifier.collision_settings.self_friction = 5  # Friction during self-collision
    cloth_modifier.collision_settings.self_distance_min = 0.015  # Minimum distance between self-collisions (tweak if needed)

    return cloth_modifier


# Function to add collision physics to object
def add_collision_physics(object):
    # Ensure the object is a mesh
    if object.type != 'MESH':
        print(f"Error: {object.name} is not a mesh object.")
        return

    # Add collision modifier
    collision_modifier = object.modifiers.new(name="Collision", type='COLLISION')
    
    # Set default collision settings (you can customize these)
    collision_modifier.settings.thickness_outer = 0.8  # Outer thickness of collision boundary
    collision_modifier.settings.thickness_inner = 0.8  # Inner thickness of collision boundary
    collision_modifier.settings.damping = 0.5  # Damping factor for collision
    collision_modifier.settings.cloth_friction = 0.5  # Friction factor for collision


# Function to add subdivision modifier
def add_subdivision_modifier(obj, levels=3):
    # Ensure the object is a mesh
    if obj.type != 'MESH':
        print(f"Error: {obj.name} is not a mesh object.")
        return
    
    # Add Cloth modifier
    subdivision_modifier = obj.modifiers.new(name="Subdivision", type='SUBSURF')
    
    # Set subdivision settings
    subdivision_modifier.levels = levels  # Viewport levels
    subdivision_modifier.render_levels = levels  # render levels


# Function makes a thin film of the mesh
def add_solidify_modifier(obj, thickness=0.001, offset=None):
    # Add a Solidify modifier for thickness
    solidify_mod = obj.modifiers.new(name="Solidify", type='SOLIDIFY')
    
    # Set the thickness
    solidify_mod.thickness = thickness
    
    # Optional: Set other solidify settings
    if offset is not None:
        solidify_mod.offset = offset
    else:
        solidify_mod.offset = -1.0
    solidify_mod.use_rim = True  # Fill in any open edges

def add_corrective_smooth_modifier(obj):
    # Ensure the object is a mesh
    if obj.type != 'MESH':
        print(f"Error: {obj.name} is not a mesh object.")
        return

    # Add a corrective smooth modifier
    csmooth_mod = obj.modifiers.new(name="CorrectiveSmooth",type='CORRECTIVE_SMOOTH')

    # Adjust parameters
    csmooth_mod.factor = 1
    csmooth_mod.iterations = 10
    csmooth_mod.scale = 0


# Function to apply auto smooth shading
def enable_auto_smooth_shading(obj, angle=30.0):
    # Ensure the object is selected and active
    bpy.context.view_layer.objects.active = obj
    obj.select_set(True)

    # Enable Auto Smooth for the object
    bpy.ops.object.shade_auto_smooth()


# Function to create a silk like material
def create_fleshy_bsdf(nodes, location, base_color):
    # Add a Principled BSDF shader node
    bsdf = nodes.new(type='ShaderNodeBsdfPrincipled')
    bsdf.location = location
    
    # Set the base color to royal red
    bsdf.inputs['Base Color'].default_value = base_color  # RGB + Alpha (Red with a bit of saturation)
    
    # Set roughness
    bsdf.inputs['Roughness'].default_value = 0.3  # Lower roughness makes the surface shiny
    
    # Set IOR
    bsdf.inputs['IOR'].default_value = 4.

    # Set metallic
    bsdf.inputs['Metallic'].default_value = 0.

    # Specular and sheen-related properties
    bsdf.inputs['Specular IOR Level'].default_value = 0.5 
    bsdf.inputs['Specular Tint'].default_value = (1, 0.1, 0.1, 1)
    
    # Sheen settings
    bsdf.inputs['Sheen Weight'].default_value = 0.8  # Moderate sheen
    bsdf.inputs['Sheen Roughness'].default_value = 0.1  # Smooth sheen
    bsdf.inputs['Sheen Tint'].default_value = (0.9, 0.05, 0.05, 1)  # Slight red tint to sheen
    
    # Subsurface settings
    bsdf.inputs['Subsurface Weight'].default_value = 0.8
    bsdf.inputs['Subsurface Radius'].default_value = (1.2, 0.8, 0.6)
    bsdf.inputs['Subsurface Scale'].default_value = 1
    bsdf.inputs['Subsurface Anisotropy'].default_value = 0.8

    return bsdf


# Function to create a silk like material
def create_brain_material():
    # Create a new material
    brain_material = bpy.data.materials.new(name="brain_material")
    
    # Enable the use of nodes (necessary for Principled BSDF)
    brain_material.use_nodes = True
    
    # Get the material's node tree
    nodes = brain_material.node_tree.nodes
    links = brain_material.node_tree.links
    
    # Clear any default nodes
    for node in nodes:
        nodes.remove(node)
    
    # Add a Principled BSDF shader node
    bsdf = create_fleshy_bsdf(nodes, location=(0, 0), base_color=(0.95, 0.3, 0.3, 1))
    bsdf2 = create_fleshy_bsdf(nodes, location=(0, -400), base_color=(0.3, 0.02, 0.02, 1))

    # Mix the two
    mixshader = nodes.new(type='ShaderNodeMixShader')
    mixshader.location = (300, 0)
    # mixshader.inputs['Fac'].default_value = 0.
    links.new(bsdf.outputs['BSDF'], mixshader.inputs[1])
    links.new(bsdf2.outputs['BSDF'], mixshader.inputs[2])

    # two color ramps for voronois
    cramp1 = nodes.new(type='ShaderNodeValToRGB')
    cramp1.location = (-300, 400)
    cramp1.color_ramp.elements[0].position = 0.008
    cramp1.color_ramp.elements[0].color = (1., 1., 1., 1.)
    cramp1.color_ramp.elements[1].position = 0.01
    cramp1.color_ramp.elements[1].color = (0., 0., 0., 1.)
    cramp2 = nodes.new(type='ShaderNodeValToRGB')
    cramp2.location = (-300, 0)
    cramp2.color_ramp.elements[0].position = 0.1
    cramp2.color_ramp.elements[0].color = (1., 1., 1., 1.)
    cramp2.color_ramp.elements[1].position = 0.04
    cramp2.color_ramp.elements[1].color = (0., 0., 0., 1.)

    # two voronois
    voronoi1 = nodes.new(type='ShaderNodeTexVoronoi')
    voronoi1.location = (-500, 400)
    voronoi1.feature = 'DISTANCE_TO_EDGE'
    voronoi1.inputs['Scale'].default_value = 3.
    voronoi1.inputs['Detail'].default_value = 2.
    links.new(voronoi1.outputs['Distance'], cramp1.inputs['Fac'])
    voronoi2 = nodes.new(type='ShaderNodeTexVoronoi')
    voronoi2.location = (-500, 0)
    voronoi2.feature = 'DISTANCE_TO_EDGE'
    voronoi2.inputs['Scale'].default_value = 30.
    links.new(voronoi2.outputs['Distance'], cramp2.inputs['Fac'])

    # Noise to wobble voronoi result
    noise = nodes.new(type='ShaderNodeTexNoise')
    noise.location = (-700, 200)
    noise.inputs['Scale'].default_value = 3.
    noise.inputs['Detail'].default_value = 15.
    links.new(noise.outputs['Color'], voronoi1.inputs['Vector'])
    links.new(noise.outputs['Color'], voronoi2.inputs['Vector'])

    # Add a math node to mix
    mixmath = nodes.new(type='ShaderNodeMath')
    mixmath.location = (0, 200)
    links.new(cramp1.outputs['Color'], mixmath.inputs[0])
    links.new(cramp2.outputs['Color'], mixmath.inputs[1])
    links.new(mixmath.outputs['Value'], mixshader.inputs['Fac'])

    # Add a displacement node
    displace = nodes.new(type='ShaderNodeDisplacement')
    displace.location = (300, 200)
    displace.inputs['Scale'].default_value = 0.01
    links.new(mixmath.outputs['Value'], displace.inputs['Height'])
    
    # Add an Output node and link it to the BSDF
    output = nodes.new(type='ShaderNodeOutputMaterial')
    output.location = (500, 0)
    links.new(mixshader.outputs['Shader'], output.inputs['Surface'])
    links.new(displace.outputs['Displacement'], output.inputs['Displacement'])

    # Let's create a bump node to improve texture
    bump = nodes.new(type='ShaderNodeBump')
    bump.location = (-200, -300)
    bump.inputs['Strength'].default_value = 0.8
    links.new(bump.outputs['Normal'], bsdf.inputs['Normal'])
    links.new(bump.outputs['Normal'], bsdf2.inputs['Normal'])
    links.new(mixmath.outputs['Value'], bump.inputs['Height'])

    # show bump and displacement
    brain_material.displacement_method = 'BOTH'

    return brain_material


# Function to create a brain like material
def create_silk_material():
    # Create a new material
    silk_material = bpy.data.materials.new(name="Royal_Red_Silk")
    
    # Enable the use of nodes (necessary for Principled BSDF)
    silk_material.use_nodes = True
    
    # Get the material's node tree
    nodes = silk_material.node_tree.nodes
    links = silk_material.node_tree.links
    
    # Clear any default nodes
    for node in nodes:
        nodes.remove(node)
    
    # Add a Principled BSDF shader node
    bsdf = nodes.new(type='ShaderNodeBsdfPrincipled')
    bsdf.location = (0, 0)
    
    # Set the base color to royal red
    bsdf.inputs['Base Color'].default_value = (0.3, 0.01, 0.03, 1)  # RGB + Alpha (Red with a bit of saturation)
    
    # Set a low roughness for silk-like smoothness
    bsdf.inputs['Roughness'].default_value = 0.5  # Lower roughness makes the surface shiny
    
    # Set a low roughness for silk-like material
    bsdf.inputs['Metallic'].default_value = 0.7

    # # Specular and sheen-related properties for silk fabric
    # bsdf.inputs['Specular IOR Level'].default_value = 0.8  # Adjust specular reflection
    # bsdf.inputs['Specular Tint'].default_value = (1, 0.1, 0.1, 1)  # White specular reflection
    
    # # Anisotropy for directional fabric-like sheen
    # bsdf.inputs['Anisotropic'].default_value = 0.8  # High anisotropy
    # bsdf.inputs['Anisotropic Rotation'].default_value = 0.2  # Low anisotropic rotation
    
    # # Sheen settings for fabric glow
    # bsdf.inputs['Sheen Weight'].default_value = 0.5  # Moderate sheen
    # bsdf.inputs['Sheen Roughness'].default_value = 0.1  # Smooth sheen
    # bsdf.inputs['Sheen Tint'].default_value = (0.9, 0.05, 0.05, 1)  # Slight red tint to sheen
    
    # Add an Output node and link it to the BSDF
    output = nodes.new(type='ShaderNodeOutputMaterial')
    output.location = (400, 0)
    links.new(bsdf.outputs['BSDF'], output.inputs['Surface'])
    
    return silk_material


# Function to create a brain like material
def create_glass_material():
    # Create a new material
    glass_material = bpy.data.materials.new(name="glass_material")
    
    # Enable the use of nodes (necessary for Principled BSDF)
    glass_material.use_nodes = True
    
    # Get the material's node tree
    nodes = glass_material.node_tree.nodes
    links = glass_material.node_tree.links
    
    # Clear any default nodes
    for node in nodes:
        nodes.remove(node)
    
    # Add a Principled BSDF shader node
    bsdf = nodes.new(type='ShaderNodeBsdfPrincipled')
    bsdf.location = (0, 0)
    
    # Set the base color to white
    bsdf.inputs['Base Color'].default_value = (1., 1., 1., 1)  # RGB + Alpha (Red with a bit of saturation)
    
    # Set roughness to control the smoothness (0 for perfect glass)
    bsdf.inputs['Roughness'].default_value = 0.0  # No roughness for clear glass

    # Set the IOR (Index of Refraction) for glass (~1.45-1.55)
    bsdf.inputs['IOR'].default_value = 1.45

    # Set the transmission value to 1 for glass-like transparency
    bsdf.inputs['Transmission Weight'].default_value = 1.0  # Full glass effect
    bsdf.inputs['Alpha'].default_value = 0.1  # Also a bit more transparent
    
    # Add an Output node and link it to the BSDF
    output = nodes.new(type='ShaderNodeOutputMaterial')
    output.location = (400, 0)
    links.new(bsdf.outputs['BSDF'], output.inputs['Surface'])
    
    return glass_material


# Function to create a brain like material
def create_colored_material(color=(1,1,1,1), name='colored_material'):
    # Create a new material
    colored_material = bpy.data.materials.new(name=name)
    
    # Enable the use of nodes (necessary for Principled BSDF)
    colored_material.use_nodes = True
    
    # Get the material's node tree
    nodes = colored_material.node_tree.nodes
    links = colored_material.node_tree.links
    
    # Clear any default nodes
    for node in nodes:
        nodes.remove(node)
    
    # Add a Principled BSDF shader node
    bsdf = nodes.new(type='ShaderNodeBsdfPrincipled')
    bsdf.location = (0, 0)
    
    # Set the base color to white
    bsdf.inputs['Base Color'].default_value = color
    
    # Set roughness
    bsdf.inputs['Roughness'].default_value = 1.0 
    
    # Add an Output node and link it to the BSDF
    output = nodes.new(type='ShaderNodeOutputMaterial')
    output.location = (400, 0)
    links.new(bsdf.outputs['BSDF'], output.inputs['Surface'])
    
    return colored_material


# Function to create a brain like material
def create_colored_light_material(color=(1,1,1,1), name='colored_material', emission_strength=1.):
    # Create a new material
    clight_material = bpy.data.materials.new(name=name)
    
    # Enable the use of nodes (necessary for Principled BSDF)
    clight_material.use_nodes = True
    
    # Get the material's node tree
    nodes = clight_material.node_tree.nodes
    links = clight_material.node_tree.links
    
    # Clear any default nodes
    for node in nodes:
        nodes.remove(node)
    
    # Add a Principled BSDF shader node
    bsdf = nodes.new(type='ShaderNodeBsdfPrincipled')
    bsdf.location = (0, 0)
    
    # Set the base color to white
    bsdf.inputs['Base Color'].default_value = color
    
    # Set roughness
    bsdf.inputs['Roughness'].default_value = 1.0 
    
    # Set emission
    bsdf.inputs['Emission Color'].default_value = color
    bsdf.inputs['Emission Strength'].default_value = emission_strength
    
    # Add an Output node and link it to the BSDF
    output = nodes.new(type='ShaderNodeOutputMaterial')
    output.location = (400, 0)
    links.new(bsdf.outputs['BSDF'], output.inputs['Surface'])
    
    return clight_material


# Function to assign material to object
def assign_material_to_object(material, obj):
    # Assign the material to the object
    if obj.data.materials:
        # If there are materials, assign to first slot
        obj.data.materials[0] = material
    else:
        # If no materials, append a new one
        obj.data.materials.append(material)


# Function to add a camera
def create_camera(camera_location, target_location):
    camera_data = bpy.data.cameras.new(name="Camera")
    camera_object = create_data_object("Camera", camera_data)
    camera_object.location = camera_location

    # Add a target (empty object) for the camera to focus on
    target_object = create_data_object("CameraTarget", None)
    target_object.location = target_location

    # Add the "Track To" constraint to the camera to make it always look at the target
    constraint = camera_object.constraints.new(type='TRACK_TO')
    constraint.target = target_object

    # parent to target to aid camera rotations
    camera_object.parent = target_object

    # Set the camera as the active camera for the scene
    bpy.context.scene.camera = camera_object

    return camera_object, target_object


# Function to add a sunlight
def create_sunlight(name, rotation, strength=1.0):
    # Add a new sun light to the scene
    light_data = bpy.data.lights.new(name=name, type='SUN')
    light_object = create_data_object(name, light_data)

    # Position the light in the scene
    light_object.rotation_euler = rotation

    # Set the strength of the light
    light_data.energy = 1.0  # You can adjust this value to make the scene brighter or dimmer
    return light_object


# Function to change background color and lighting
def set_background(background_color, ambient_light_strength=1.0):
    bpy.context.scene.world.use_nodes = True
    bg = bpy.context.scene.world.node_tree.nodes.get("Background")
    bg.inputs[0].default_value = background_color
    bg.inputs[1].default_value = ambient_light_strength  # Strength of the ambient light


# Function to set HDRI background
def set_hdri_background(hdr_image_path):
    # HDRI background
    bpy.context.scene.world.use_nodes = True
    node_tree = bpy.context.scene.world.node_tree
    # Clear any existing nodes in the world
    nodes = node_tree.nodes
    nodes.clear()
    # Add Background node and set it as the output
    background_node = nodes.new(type="ShaderNodeBackground")
    # Add Environment Texture node (for the HDRI)
    env_texture_node = nodes.new(type="ShaderNodeTexEnvironment")
    # Load your HDRI image here (replace with your file path)
    env_texture_node.image = bpy.data.images.load(hdr_image_path)
    # Add World Output node
    world_output_node = nodes.new(type="ShaderNodeOutputWorld")
    # Connect the Environment Texture node to the Background node
    node_tree.links.new(env_texture_node.outputs["Color"], background_node.inputs["Color"])
    # Connect the Background node to the World Output node
    node_tree.links.new(background_node.outputs["Background"], world_output_node.inputs["Surface"])
    # Optional: Adjust the strength of the HDRI light
    background_node.inputs["Strength"].default_value = 1.5  # You can change this value



# Function to set the resolution of renders
def set_render_properties(resolution_x=1920, resolution_y=1080, resolution_percentage=100, fps=24):
    bpy.context.scene.render.resolution_x = resolution_x  # Width in pixels
    bpy.context.scene.render.resolution_y = resolution_y  # Height in pixels
    bpy.context.scene.render.resolution_percentage = resolution_percentage  # Scale
    bpy.context.scene.render.fps = fps # frame rate


# Render a single frame to a PNG
def render_single_frame(frame, filepath, engine='BLENDER_EEVEE_NEXT', transparent=True):
    # Set the current frame to render
    bpy.context.scene.frame_set(frame)

    # Set render output file path
    bpy.context.scene.render.filepath = filepath

    # Set render engine to 'CYCLES' or 'BLENDER_EEVEE_NEXT'
    bpy.context.scene.render.engine = engine
    
    # Set the file format to PNG
    bpy.context.scene.render.image_settings.file_format = 'PNG'

    # Make the background transparent
    if transparent:
        bpy.context.scene.render.film_transparent = True
        bpy.context.scene.render.image_settings.color_mode = 'RGBA'  # Enable alpha (transparency)

    # Trigger the render and save the image
    bpy.ops.render.render(write_still=True)


# Render an animation to MP4
def render_animation(filepath, engine='BLENDER_EEVEE_NEXT', render=False):
    # Set render output file path
    bpy.context.scene.render.filepath = filepath

    # Set render engine to 'CYCLES' or 'BLENDER_EEVEE_NEXT'
    bpy.context.scene.render.engine = engine
    bpy.context.scene.render.use_persistent_data = True  # Added to speedup rendering
    if engine=='CYCLES':
        bpy.context.scene.cycles.adaptive_threshold = 0.1 # Added to speedup rendering

    # Set the file format to FFmpeg video
    bpy.context.scene.render.image_settings.file_format = 'FFMPEG'
    
    # Set the codec and encoding options
    bpy.context.scene.render.ffmpeg.format = 'MPEG4'
    bpy.context.scene.render.ffmpeg.codec = 'H264'
    bpy.context.scene.render.ffmpeg.constant_rate_factor = 'HIGH'

    # Render the animation
    if render:
        bpy.ops.render.render(animation=True)


# Render an animation with transparency to WEBM
def render_transparent_animation(filepath, engine='BLENDER_EEVEE_NEXT'):
    # Set render output file path
    bpy.context.scene.render.filepath = filepath

    # Set render engine to 'CYCLES' or 'BLENDER_EEVEE_NEXT'
    bpy.context.scene.render.engine = engine
    bpy.context.scene.render.use_persistent_data = True  # Added to speedup rendering
    if engine=='CYCLES':
        bpy.context.scene.cycles.adaptive_threshold = 0.1 # Added to speedup rendering
    
    # Enable transparency for the world background
    bpy.context.scene.render.film_transparent = True
    
    # Set the file format to FFmpeg video
    bpy.context.scene.render.image_settings.file_format = 'FFMPEG'
    
    # Set the codec and encoding options
    bpy.context.scene.render.ffmpeg.format = 'WEBM'  # or 'QUICKTIME' for .mov
    bpy.context.scene.render.ffmpeg.codec = 'VP9'    # 'WEBM' or 'QTRLE' for QuickTime
    
    # Ensure the color mode is RGBA (for transparency)
    bpy.context.scene.render.image_settings.color_mode = 'RGBA'

    # Render the animation
    bpy.ops.render.render(animation=True)

# loading brain data
################################################################################
################################################################################
(
    left_vertices_loaded,
    left_triangles_loaded,
    right_vertices_loaded,
    right_triangles_loaded
) = load_cortical_meshes('inflated')

# Brain scaling (original is mm coordinates)
brain_scale = 50
left_vertices_loaded /= (1000 / brain_scale)
right_vertices_loaded /= (1000 / brain_scale)

dscalar = nib.load(cbu.cifti_template_file)
brain_models = [x for x in dscalar.header.get_index_map(1).brain_models]
left_cortical_surface_model, right_cortical_surface_model = brain_models[0], brain_models[1]

gradients = nib.load("/mnt/local_storage/Research/Codes/fMRI/DataStore/Templates/principal_gradient/hcp.gradients.dscalar.nii")

# read subcortex atlas files
subcortical_mesh_infos = {}
subcortical_atlas_infos = {}
msa_scales = ['1', '2', '3', '4']
for msa_scale in msa_scales:
    msa = nib.load(f"/mnt/local_storage/Research/Codes/fMRI/DataStore/Templates/MSA/Tian_Subcortex_S{msa_scale}_3T_1mm.nii.gz")
    msa_info = read_msa_atlas_information(f"/mnt/local_storage/Research/Codes/fMRI/DataStore/Templates/MSA/Schaefer2018_100Parcels_7Networks_order_Tian_Subcortex_S{msa_scale}_label.txt")
    subcortical_atlas_infos[msa_scale] = msa_info
    subcortical_mesh_infos[msa_scale] = {}
    for index in np.unique(msa.get_fdata()):
        if index != 0:
            voxels_ijk = np.array(np.where(msa.get_fdata() == index)).T
            transformation_matrix = msa.affine
            surface_vertices, surface_triangles = cbu.generate_surface_marching_cube(
                voxels_ijk, transformation_matrix, smoothing_filter="laplacian",
                simplify=False, smoothing=0, subdivide=False,
            )
            surface_vertices /= (1000 / brain_scale)
            subcortical_mesh_infos[msa_scale][index] = (surface_vertices, surface_triangles)

# Blender scripting starts here
################################################################################
################################################################################
# First clean the scene
blender_clean_up()
################################################################################


################################################################################
################################################################################
# Frame timing
start_frame = 1
end_frame = 600
frame_setup(start_frame=start_frame, end_frame=end_frame)
################################################################################

################################################################################
################################################################################
# Now create the mesh
base_offset = (left_vertices_loaded.mean(0) + right_vertices_loaded.mean(0))/2
left_cortex_mesh, left_cortex_obj = create_mesh_object(
    left_vertices_loaded - base_offset, left_triangles_loaded,
    mesh_name="left_cortex_mesh", object_name="left_cortex_obj"
)
right_cortex_mesh, right_cortex_obj = create_mesh_object(
    right_vertices_loaded - base_offset, right_triangles_loaded,
    mesh_name="right_cortex_mesh", object_name="right_cortex_obj"
)
# and the subcortical regions
subcortical_meshes = {}
subcortical_objs = {}
for msa_scale in msa_scales:
    msa_info = subcortical_atlas_infos[msa_scale]
    for index in subcortical_mesh_infos[msa_scale]:
        subcortical_meshes[f"msa_s{msa_scale}_p{index}"], subcortical_objs[f"msa_s{msa_scale}_p{index}"] = create_mesh_object(
            subcortical_mesh_infos[msa_scale][index][0] - base_offset,
            subcortical_mesh_infos[msa_scale][index][1],
            mesh_name=f"msa_s{msa_scale}_p{index}_{msa_info[index]['name']}_mesh",
            object_name=f"msa_s{msa_scale}_p{index}_{msa_info[index]['name']}_obj"
        )
        # smooth mesh
        add_corrective_smooth_modifier(subcortical_objs[f"msa_s{msa_scale}_p{index}"])
################################################################################

################################################################################
################################################################################
# Apply smooth shading
shade_smooth(left_cortex_obj)
shade_smooth(right_cortex_obj)
for key in subcortical_objs:
    shade_smooth(subcortical_objs[key])
################################################################################

################################################################################
################################################################################
# handling colors
# A simple glass material
glass_material = create_glass_material()

# Make the brain look more real
assign_material_to_object(glass_material, left_cortex_obj)
assign_material_to_object(glass_material, right_cortex_obj)

# Make the glass thin to avoid strange transmissions
add_solidify_modifier(left_cortex_obj, thickness=0.0001)
add_solidify_modifier(right_cortex_obj, thickness=0.0001)

# colored subcortex material
sc_materials = {}
for msa_scale in msa_scales:
    msa_info = subcortical_atlas_infos[msa_scale]
    for index in subcortical_mesh_infos[msa_scale]:
        sc_materials[f"msa_s{msa_scale}_p{index}"] = create_colored_light_material(
            color=(msa_info[index]['r'], msa_info[index]['g'], msa_info[index]['b'],1),
            name=f"msa_s{msa_scale}_p{index}_{msa_info[index]['name']}_material", emission_strength=0.1)
        assign_material_to_object(sc_materials[f"msa_s{msa_scale}_p{index}"], subcortical_objs[f"msa_s{msa_scale}_p{index}"])
################################################################################


################################################################################
################################################################################
# handle parents and groups
cortex_object = create_data_object("cortex", None)
cortex_object.location = (0, 0, 0)
left_cortex_obj.parent = cortex_object
right_cortex_obj.parent = cortex_object
msa_objects = {}
for msa_scale in msa_scales:
    msa_objects[msa_scale] = create_data_object(f"msa_s{msa_scale}", None)
    msa_objects[msa_scale].location = (0, 0, 0)
    for index in subcortical_mesh_infos[msa_scale]:
        subcortical_objs[f"msa_s{msa_scale}_p{index}"].parent = msa_objects[msa_scale]


# Add a fixed pivot to move all brain parts together
pivot_object = create_data_object("BrainPivot", None)
pivot_object.location = (0, 0, 0)
# pivot_object.location = (left_vertices_loaded.mean(0) + right_vertices_loaded.mean(0))/2

# parent to pivot to aid brain rotations
cortex_object.parent = pivot_object
for msa_scale in msa_scales:
    msa_objects[msa_scale].parent = pivot_object
################################################################################

################################################################################
################################################################################
# Add texts on top of view
font_path = "/home/sina/.fonts/Ubuntu Bold Italic Nerd Font Complete.ttf"
colors = plt.cm.get_cmap('rainbow', len(msa_scales))
font_colors = [colors(i) for i in range(len(msa_scales))]
msa_texts = {}
msa_text_materials = {}
for i, msa_scale in enumerate(msa_scales):
    msa_text_materials[msa_scale] = create_colored_material(
        color=font_colors[i], name="text_color_material"
    )
    msa_texts[msa_scale] = create_text_object(
        name=f"msa_s{msa_scale}_text", location=(0, 0, 40),
        rotation=(math.radians(90), 0, math.radians(-90)),
        text=f"MSA: scale {msa_scale}", align_x="CENTER",
        align_y="TOP", font=font_path,
    )
    add_solidify_modifier(msa_texts[msa_scale], thickness=0.1)
    assign_material_to_object(msa_text_materials[msa_scale], msa_texts[msa_scale])
################################################################################


################################################################################
################################################################################
# Create a wall behind view
wall = create_plane(
    name="WallPlane", scale=(50, 50),
    location=(
        left_vertices_loaded[:,0].max() + 20,
        left_vertices_loaded[:,1].mean(),
        left_vertices_loaded[:,2].mean()
    ), rotation=(0,math.radians(90),0), subdivision_kws={'subdivide': True, 'cuts': 20}
)
wall_material = create_colored_material(color=(0.,0.,0.,1), name="wall_material")
assign_material_to_object(wall_material, wall)
################################################################################


################################################################################
################################################################################
# Add a camera object
camera_object, target_object = create_camera(
    camera_location=brain_scale * np.array((-0.3, 0, 0)),
    target_location=(0,0,0)
    # target_location=(left_vertices_loaded.mean(0) + right_vertices_loaded.mean(0))/2
)

# camera animation
for i in range(end_frame):
    # Keyframe the rotation of the empty (animate around the Z axis)
    target_object.rotation_euler = (
        0,
        np.cos(2 * np.pi * i / end_frame)**3 * math.radians(4),
        np.sin(4 * np.pi * i / end_frame) * math.radians(2)
    )
    target_object.keyframe_insert(data_path="rotation_euler", frame=i + 1)
################################################################################


################################################################################
################################################################################
# # add multiple surrounding light sources
light_count = 6
# for i in range(light_count + 1):
#     # Add a new sun light to the scene
#     create_sunlight(
#         name=f"SunLight_{i}",
#         rotation=(np.pi / 2, 0, (-np.pi * (i / light_count)) - (np.pi/2)),
#         strength=1.0
#     )
# and a light source from above
create_sunlight(
    name=f"SunLight_{light_count + 1}",
    rotation=(0, 0, 0),
    strength=2.0
)
################################################################################

# Function to appear and disappear by animation
def appearance_animation(material, appear_start, appear_end, disappear_start, disappear_end):
    bsdf = material.node_tree.nodes["Principled BSDF"]
    bsdf.inputs["Alpha"].default_value = 0
    bsdf.inputs["Alpha"].keyframe_insert(data_path='default_value', frame=appear_start)
    bsdf.inputs["Alpha"].default_value = 1
    bsdf.inputs["Alpha"].keyframe_insert(data_path='default_value', frame=appear_end)
    bsdf.inputs["Alpha"].default_value = 1
    bsdf.inputs["Alpha"].keyframe_insert(data_path='default_value', frame=disappear_start)
    bsdf.inputs["Alpha"].default_value = 0
    bsdf.inputs["Alpha"].keyframe_insert(data_path='default_value', frame=disappear_end)

# Function to bounce text in
def obj_appearance_animation(obj, loc0, loc1, appear_start, appear_end):
    obj.location = loc0
    obj.keyframe_insert(data_path='location', frame=appear_start - 1)
    obj.location = loc1
    obj.keyframe_insert(data_path='location', frame=appear_start)
    obj.location = loc1
    obj.keyframe_insert(data_path='location', frame=appear_end)
    obj.location = loc0
    obj.keyframe_insert(data_path='location', frame=appear_end + 1)

################################################################################
################################################################################
# Animation directives

# brain rotations

# Start rotation
pivot_object.rotation_euler = (0., 0., 0.)
pivot_object.keyframe_insert(data_path="rotation_euler", frame=start_frame)

# End rotation
pivot_object.rotation_euler = (0., 0., 8*np.pi)
pivot_object.keyframe_insert(data_path="rotation_euler", frame=end_frame)

# Atlas animations
for i, msa_scale in enumerate(msa_scales):
    t = end_frame // len(msa_scales)
    times = (t * i) + 5, (t * i) + 15, (t * i) + t - 15, (t * i) + t - 5
    appearance_animation(msa_text_materials[msa_scale], *times)
    obj_appearance_animation(msa_texts[msa_scale], (0, 0, 40), (0, 0, 4), (t * i), (t * i) + t)
    obj_appearance_animation(msa_objects[msa_scale], (0, 0, -40), (0, 0, 0), (t * i), (t * i) + t)
    for index in subcortical_mesh_infos[msa_scale]:
        appearance_animation(sc_materials[f"msa_s{msa_scale}_p{index}"], *times)

################################################################################

print("here")

################################################################################
################################################################################
# Set the world background color
set_background(background_color=(1, 1, 1, 1), ambient_light_strength=1.2)
# Set the world background from hdri
# hdr_image_path = "/mnt/local_storage/Research/Codes/fMRI/DataStore/Templates/HDRI/evening_road_01_puresky_8k.hdr"
# set_hdri_background(hdr_image_path)
################################################################################

################################################################################
################################################################################
# Set render resolution
set_render_properties(resolution_x=1000, resolution_y=800, fps=30)
################################################################################

################################################################################
################################################################################
# render a single frame and store as image
# render_single_frame(frame=1, filepath="/mnt/local_storage/Research/Codes/fMRI/DataStore/tmp/blender/rendered_image.png", engine='BLENDER_EEVEE_NEXT')
################################################################################

################################################################################
################################################################################
# render animation
render_animation(filepath="/mnt/local_storage/Research/Codes/fMRI/DataStore/tmp/blender/glass_brain_msa.mp4", engine='BLENDER_EEVEE_NEXT')
################################################################################

################################################################################
################################################################################
# bpy.ops.screen.animation_play()
################################################################################


################################################################################
################################################################################
# # Finally save the project
# # Specify the file path where you want to save the Blender project
# file_path = "/mnt/local_storage/Research/Codes/fMRI/DataStore/tmp/blender/project_cloth.blend"

# # Save the current Blender file
# bpy.ops.wm.save_as_mainfile(filepath=file_path)
################################################################################


bpy.context.scene.render.engine = 'CYCLES'
# bpy.context.space_data.shading.type = 'RENDERED'

This is the resulting render!

glass_brain_msa_final.mp4