Download generate_brain_data.py from 40Hz/brain-viz: direct link, hf CLI and curl.
- Browser
- Download file 12.6 kB
-
https://huggingface.co/spaces/40Hz/brain-viz/resolve/main/generate_brain_data.py
- Command line
-
hf download hf://spaces/40Hz/brain-viz/generate_brain_data.py
-
curl -L -o generate_brain_data.py https://huggingface.co/spaces/40Hz/brain-viz/resolve/main/generate_brain_data.py
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() | |