vertexColors_v1.py

vertexColors_v1.py

"""
vertexColors_v1.py

@author: wronk

Script meant to be run within blender to load source estimate data and create a series of images
(one for each time point) that can be converted into a movie. This requires that the brain surface
first be loaded into blender. 
"""
# Blender specific libraries:
import bpy
from mathutils import Color, Vector

# Common libraries
import numpy as np
import random
import os.path as op

data_dir = '/media/Toshiba/Blender/Ross_AudVis/'
save_folder = 'images'
subject = 'fsaverage'
hemi = 'lh'
f_names = ['Visual', 'Target', 'Masker']
brain_str_key = 'fsaverage.lh.inflated_pre'
vert_list = []
data_list = []
frames = [20, 540]
control_pts = [2, 2.5, 6]  # Control points for 3 color sliders
max_lim = np.max(control_pts)
control_pts /= max_lim

cam = bpy.data.objects['Camera']
brain_mat = 'Material.001'
colorRamp = bpy.data.materials[brain_mat].node_tree.nodes["ColorRamp"]
#origin = bpy.data.objects['Empty']
bpy.ops.mesh.vertex_color_remove()

def color_rescale(value):
    return (value, value, value)

for fi, f_name in enumerate(f_names):
    temp_verts = np.load(op.join(data_dir, 'stcs/vertices', f_name + '-' +
                                 hemi + '.npy'))
    vert_list.append(temp_verts)
    temp_data = np.load(op.join(data_dir, 'stcs/data', f_name + '-' +
                                hemi + '.npy'))[:, frames[0]:frames[1]]
    # XXX Taking abs of z-scores and dividing by max of control_pts
    data_list.append(np.abs(temp_data) / max_lim)
    
color_map = None
for file_i, file_name in enumerate(f_names):
    for cpi, cp in enumerate(control_pts):
        colorRamp.color_ramp.elements[cpi].position = cp
    for frame_i in range(frames[1] - frames[0]):
        my_object = bpy.data.objects[brain_str_key].data
        vert_list = my_object.vertices
        if color_map is None:
            color_map = my_object.vertex_colors.new()
        c_len = len(my_object.vertex_colors[0].data)
        for poly in my_object.polygons:
            for loop_i in poly.loop_indices:
                loop = my_object.loops[loop_i]
                v = loop.vertex_index
                col = Color(color_rescale(data_list[file_i][v, frame_i]))
                color_map.data[loop_i].color = col
        bpy.data.scenes["Scene"].render.filepath = \
            op.join(data_dir, save_folder, file_name, subject + '_' + hemi +
                    '_' + file_name + '_%d.jpg' % (frame_i + frames[0]))
        bpy.ops.render.render(write_still=True)