brain-viz / generate_brain_data.py
40Hz's picture
Upload generate_brain_data.py with huggingface_hub
96f52d3 verified
Raw History Blame Contribute Delete
12.6 kB
#!/usr/bin/env python3
"""Generate brain-mesh.json and state-data.json for Three.js brain visualizer.
If nilearn is available, uses fsaverage5 mesh + Destrieux atlas.
Otherwise generates a simplified brain-shaped mesh using subdivided icosahedron.
"""
import json, math, os, sys
try:
import numpy as np
HAS_NUMPY = True
except ImportError:
np = None
HAS_NUMPY = False
try:
from nilearn import datasets, surface
HAS_NILEARN = True
except ImportError:
HAS_NILEARN = False
# ─── Icosahedron subdivision ───────────────────────────────────────────────
G = (1 + math.sqrt(5)) / 2
ICOSA_VERTICES = [
(-1, G, 0), (1, G, 0), (-1, -G, 0), (1, -G, 0),
(0, -1, G), (0, 1, G), (0, -1, -G), (0, 1, -G),
(G, 0, -1), (G, 0, 1), (-G, 0, -1), (-G, 0, 1),
]
ICOSA_FACES = [
(0, 11, 5), (0, 5, 1), (0, 1, 7), (0, 7, 10), (0, 10, 11),
(1, 5, 9), (5, 11, 4), (11, 10, 2), (10, 7, 6), (7, 1, 8),
(3, 9, 4), (3, 4, 2), (3, 2, 6), (3, 6, 8), (3, 8, 9),
(4, 9, 5), (2, 4, 11), (6, 2, 10), (8, 6, 7), (9, 8, 1),
]
def _midpoint(v1, v2, verts):
x = (verts[v1][0] + verts[v2][0]) / 2
y = (verts[v1][1] + verts[v2][1]) / 2
z = (verts[v1][2] + verts[v2][2]) / 2
length = math.sqrt(x*x + y*y + z*z)
return (x/length, y/length, z/length)
def subdivide_icosahedron(detail):
verts = list(ICOSA_VERTICES)
faces = list(ICOSA_FACES)
for _ in range(detail):
mid_cache = {}
new_faces = []
for (a, b, c) in faces:
ab = mid_cache.get((a, b)) or mid_cache.setdefault((a, b),
len(verts) + len([k for k in mid_cache if k < (a, b)]))
# We need a proper approach
# Actually, let me use a simpler index tracking approach
# Fall through to simpler implementation below
def generate_icosahedron_mesh(detail=4):
"""Generate subdivided icosahedron. detail=4 -> 2562 verts, detail=5 -> 10242."""
verts = [(x/math.sqrt(x*x+y*y+z*z), y/math.sqrt(x*x+y*y+z*z), z/math.sqrt(x*x+y*y+z*z))
for (x, y, z) in ICOSA_VERTICES]
faces = list(ICOSA_FACES)
for _ in range(detail):
midpoint_map = {}
new_faces = []
for a, b, c in faces:
key_ab = (a, b) if a < b else (b, a)
key_bc = (b, c) if b < c else (c, b)
key_ca = (c, a) if c < a else (a, c)
if key_ab not in midpoint_map:
midpoint_map[key_ab] = len(verts)
verts.append(_midpoint(key_ab[0], key_ab[1], verts))
if key_bc not in midpoint_map:
midpoint_map[key_bc] = len(verts)
verts.append(_midpoint(key_bc[0], key_bc[1], verts))
if key_ca not in midpoint_map:
midpoint_map[key_ca] = len(verts)
verts.append(_midpoint(key_ca[0], key_ca[1], verts))
ab = midpoint_map[key_ab]
bc = midpoint_map[key_bc]
ca = midpoint_map[key_ca]
new_faces.append((a, ab, ca))
new_faces.append((b, bc, ab))
new_faces.append((c, ca, bc))
new_faces.append((ab, bc, ca))
faces = new_faces
# Brain-like deformation
deformed = []
for (x, y, z) in verts:
nx, ny, nz = x, y, z
bump = 1 + 0.08 * math.sin(5*nx) * math.cos(4*ny) * math.sin(6*nz) \
+ 0.05 * math.cos(8*nx*nx) * math.sin(7*ny*ny) \
+ 0.03 * (0.5 + 0.5 * math.sin(10 * nz))
front_stretch = 1 + 0.15 * (nx * nx if nx > 0 else 0)
flat_bottom = 1 - 0.1 * max(0, -ny)
scale = bump * front_stretch * flat_bottom
deformed.append((x * scale, y * scale, z * scale))
# Add fissure (split between hemispheres)
# Push vertices near x=0, z>0 slightly apart
fissured = []
for (x, y, z) in deformed:
if abs(x) < 0.15 and z > 0.1:
x += 0.08 * (1 if z > 0 else -1)
fissured.append((x, y, z))
return {
'vertices': fissured,
'faces': faces,
}
def load_nilearn_mesh():
"""Load fsaverage5 mesh using nilearn."""
print("Loading fsaverage5 mesh via nilearn...")
fsa = datasets.fetch_surf_fsaverage('fsaverage5')
verts_l, faces_l = surface.load_surf_mesh(fsa['pial_left'])
verts_r, faces_r = surface.load_surf_mesh(fsa['pial_right'])
# Combine hemispheres
verts = np.concatenate([verts_l, verts_r], axis=0)
offset = verts_l.shape[0]
faces = np.concatenate([faces_l, faces_r + offset], axis=0)
return {
'vertices': verts.tolist(),
'faces': faces.tolist(),
}
# ─── Region definitions ─────────────────────────────────────────────────────
# For simplified state data, define ROI functions on unit sphere
def _roi_weights(x, y, z):
"""Return dict of region weights for a vertex on unit sphere."""
nx, ny, nz = x, y, z
pfc = max(0, nx)**2
parietal = max(0, ny)**2 * (1 - max(0, nx)**2)
temporal = abs(nz)**1.5 * (1 - abs(nx))
occipital = max(0, -nx)**2
dmn = (0.3 * max(0, nx)**2 + 0.3 * max(0, -nx)**2 + 0.2 * abs(nz)**2)
limbic = abs(nz)**2 * max(0, -ny + 0.2)**2
return {'pfc': pfc, 'parietal': parietal, 'temporal': temporal,
'occipital': occipital, 'dmn': dmn, 'limbic': limbic}
def generate_state_data_simple(verts_list):
"""Generate 4-state data using simplified ROI mapping."""
# Use deterministic noise based on vertex index
def _noise(i):
return 0.5 + 0.5 * math.sin(i * 123.456 + 789.012)
n = len(verts_list)
states = {'baseline': [], 'focus12': [], 'focus15': [], 'trauma': []}
for i, (x, y, z) in enumerate(verts_list):
ln = math.sqrt(x*x + y*y + z*z)
if ln == 0:
w = {k: 0 for k in ('pfc','parietal','temporal','occipital','dmn','limbic')}
else:
w = _roi_weights(x/ln, y/ln, z/ln)
nse = _noise(i)
# Baseline: centered at 0.3 with noise, DMN slightly boosted
v_base = max(0, min(1, 0.25 + 0.12 * nse + 0.1 * w['dmn']))
# Focus 12: high PFC + parietal
v_f12 = max(0, min(1, 0.1 + 0.7 * w['pfc'] + 0.5 * w['parietal'] + 0.1 * nse))
# Focus 15: very high PFC + temporal + occipital
v_f15 = max(0, min(1, 0.1 + 0.85 * w['pfc'] + 0.6 * w['temporal']
+ 0.5 * w['occipital'] + 0.08 * nse))
# Trauma: high limbic, low PFC, suppressed DMN
v_tr = max(0, min(1, 0.15 + 0.7 * w['limbic']
+ 0.05 * (1 - w['pfc'])
+ 0.05 * w['dmn']
+ 0.05 * nse))
states['baseline'].append(v_base)
states['focus12'].append(v_f12)
states['focus15'].append(v_f15)
states['trauma'].append(v_tr)
return states
def generate_state_data_nilearn(verts_list):
"""Generate state data using Destrieux atlas regions via nilearn."""
print("Generating state data with Destrieux atlas...")
try:
from nilearn import datasets, surface
import numpy as np
except ImportError:
print(" nilearn import failed, falling back to simple mapping")
return generate_state_data_simple(verts_list)
verts = np.array(verts_list)
n = verts.shape[0]
try:
destrieux = datasets.fetch_atlas_surf_destrieux()
labels = destrieux['map_left']
# Combine with right hemisphere
labels = np.concatenate([labels, destrieux['map_right']])
label_names = destrieux['labels']
except Exception as e:
print(f" Destrieux atlas failed ({e}), using simple mapping")
return generate_state_data_simple(verts_list)
# Map Destrieux labels to functional groups
# These are approximate mappings
pfc_labels = set()
parietal_labels = set()
temporal_labels = set()
occipital_labels = set()
limbic_labels = set()
dmn_labels = set()
for idx, name in label_names:
name_str = str(name).lower()
if any(k in name_str for k in ['front', 'prefront', 'precent', 'fronto']):
pfc_labels.add(idx)
if any(k in name_str for k in ['pariet', 'supramarg', 'postcent']):
parietal_labels.add(idx)
if any(k in name_str for k in ['tempor', 'fusiform', 'parahippoc']):
temporal_labels.add(idx)
if any(k in name_str for k in ['occip', 'lingual', 'cuneus', 'calcar']):
occipital_labels.add(idx)
if any(k in name_str for k in ['amygdal', 'hippocamp', 'cingul', 'insul',
'orbit', 'parolfact']):
limbic_labels.add(idx)
if any(k in name_str for k in ['cingul', 'precuneus', 'angular',
'frontopol', 'temporopol']):
dmn_labels.add(idx)
# For vertices not in any known region, use fallback
states = {'baseline': [], 'focus12': [], 'focus15': [], 'trauma': []}
for i in range(n):
lab = labels[i] if i < len(labels) else 0
in_pfc = 1.0 if lab in pfc_labels else 0.0
in_parietal = 1.0 if lab in parietal_labels else 0.0
in_temporal = 1.0 if lab in temporal_labels else 0.0
in_occipital = 1.0 if lab in occipital_labels else 0.0
in_limbic = 1.0 if lab in limbic_labels else 0.0
in_dmn = 1.0 if lab in dmn_labels else 0.0
nse = 0.5 + 0.5 * math.sin(i * 123.456 + 789.012)
# Add smooth falloff at region boundaries using spatial distance
x, y, z = verts[i]
ln = math.sqrt(x*x + y*y + z*z) or 1
nx, ny, nz = x/ln, y/ln, z/ln
spatial = _roi_weights(nx, ny, nz)
# Blend label-based with spatial
pfc_w = max(in_pfc, spatial['pfc'] * 0.3)
par_w = max(in_parietal, spatial['parietal'] * 0.3)
tmp_w = max(in_temporal, spatial['temporal'] * 0.3)
occ_w = max(in_occipital, spatial['occipital'] * 0.3)
lim_w = max(in_limbic, spatial['limbic'] * 0.3)
dmn_w = max(in_dmn, spatial['dmn'] * 0.3)
v_base = max(0, min(1, 0.25 + 0.12 * nse + 0.1 * dmn_w))
v_f12 = max(0, min(1, 0.1 + 0.7 * pfc_w + 0.5 * par_w + 0.1 * nse))
v_f15 = max(0, min(1, 0.1 + 0.85 * pfc_w + 0.6 * tmp_w + 0.5 * occ_w + 0.08 * nse))
v_tr = max(0, min(1, 0.15 + 0.7 * lim_w + 0.05 * (1 - pfc_w) + 0.05 * dmn_w + 0.05 * nse))
states['baseline'].append(v_base)
states['focus12'].append(v_f12)
states['focus15'].append(v_f15)
states['trauma'].append(v_tr)
return states
# ─── Main ───────────────────────────────────────────────────────────────────
def main():
out_dir = os.path.dirname(os.path.abspath(__file__))
# Generate mesh
if HAS_NILEARN:
try:
mesh = load_nilearn_mesh()
print(f"Loaded fsaverage5: {len(mesh['vertices'])} verts, {len(mesh['faces'])} faces")
except Exception as e:
print(f"nilearn mesh failed ({e}), using icosahedron fallback")
mesh = generate_icosahedron_mesh(detail=4)
else:
print("nilearn not available. Generating simplified brain mesh.")
mesh = generate_icosahedron_mesh(detail=4)
mesh_path = os.path.join(out_dir, 'brain-mesh.json')
with open(mesh_path, 'w') as f:
json.dump({'vertices': mesh['vertices'], 'faces': mesh['faces']}, f)
mesh_size = os.path.getsize(mesh_path)
print(f" brain-mesh.json: {mesh_size} bytes ({len(mesh['vertices'])} verts)")
# Generate state data
if HAS_NILEARN:
state_data = generate_state_data_nilearn(mesh['vertices'])
else:
state_data = generate_state_data_simple(mesh['vertices'])
# Round to 4 decimal places to reduce file size
for state in state_data:
state_data[state] = [round(v, 4) for v in state_data[state]]
state_path = os.path.join(out_dir, 'state-data.json')
with open(state_path, 'w') as f:
json.dump(state_data, f)
state_size = os.path.getsize(state_path)
print(f" state-data.json: {state_size} bytes ({len(state_data['baseline'])} vals per state)")
print(f"\nDone. Output in: {out_dir}")
print(json.dumps({
'brain-mesh.json': mesh_size,
'state-data.json': state_size,
}))
if __name__ == '__main__':
main()