Upload brain_virality_predictor/tribe_wrapper.py with huggingface_hub

brain_virality_predictor/tribe_wrapper.py

@@ -0,0 +1,204 @@
+"""
+TRIBE v2 video → brain response inference with Yeo network masking.
+Maps cortical vertices (fsaverage5) into 7 Yeo functional networks.
+"""
+
+import json, os, warnings
+from pathlib import Path
+from typing import Dict, List, Optional, Tuple
+import numpy as np
+
+try:
+    import torch
+    HAS_TORCH = True
+except ImportError:
+    HAS_TORCH = False
+
+# ── Yeo 2011 7-network fsaverage5 labels (left hemi → vertex IDs) ─────
+# These are the canonical surface-based labels from FreeSurfer's fsaverage5.
+# We ship a lightweight JSON with vertex→network mapping so nilearn is optional.
+YEO7_NAMES = {
+    1: "Visual",
+    2: "Somatomotor",
+    3: "Dorsal_Attention",
+    4: "Ventral_Attention",
+    5: "Limbic",
+    6: "Frontoparietal",
+    7: "Default",
+    0: "None",
+}
+
+# Yeo7 → our 8 UX signals (multi-mapping allowed)
+YEO7_TO_UX = {
+    "Visual":           ["aesthetic_appeal", "visual_fluency"],
+    "Somatomotor":      ["motor_readiness"],
+    "Dorsal_Attention": ["visual_fluency", "cognitive_load"],
+    "Ventral_Attention": ["surprise_novelty"],
+    "Limbic":           ["reward_anticipation", "trust_affinity"],
+    "Frontoparietal":   ["cognitive_load", "friction_anxiety"],
+    "Default":          ["trust_affinity", "aesthetic_appeal"],
+}
+
+
+class YeoAtlas:
+    """Lightweight Yeo-2011 7-network atlas for fsaverage5 surface."""
+
+    def __init__(self, label_path: Optional[str] = None):
+        self.vertex_to_network: Dict[int, int] = {}
+        self.network_vertices: Dict[str, List[int]] = {n: [] for n in YEO7_NAMES.values()}
+
+        if label_path and Path(label_path).exists():
+            self._load_label_file(label_path)
+        else:
+            self._generate_fallback()
+
+    def _load_label_file(self, path: str):
+        # FreeSurfer .label or .annot format not needed if we use our fallback JSON
+        with open(path) as f:
+            data = json.load(f)
+        self.vertex_to_network = {int(k): int(v) for k, v in data.items()}
+        for v, n in self.vertex_to_network.items():
+            name = YEO7_NAMES.get(n, "None")
+            self.network_vertices.setdefault(name, []).append(v)
+
+    def _generate_fallback(self):
+        """
+        Fallback: generate a deterministic pseudo-Yeo parcellation based on
+        standard FreeSurfer fsaverage5 vertex ordering.  This is NOT the real
+        atlas, but is sufficient to make the pipeline run if the user hasn't
+        downloaded the true labels.  We emit a loud warning.
+        """
+        warnings.warn(
+            "Yeo atlas labels not found — using deterministic fallback. "
+            "For real neuroscience, download fsaverage5 Yeo2011_7Networks. "
+            "labels from https://github.com/ThomasYeoLab/CBIG/tree/master/stable_projects/brain_parcellation/Yeo2011/"
+        )
+        n_vertices = 20484  # fsaverage5 total
+        # deterministic round-robin assignment per canonical ordering
+        for v in range(n_vertices):
+            net_id = (v % 7) + 1
+            self.vertex_to_network[v] = net_id
+            name = YEO7_NAMES[net_id]
+            self.network_vertices[name].append(v)
+
+    def mask_for_network(self, name: str) -> np.ndarray:
+        verts = np.array(self.network_vertices.get(name, []), dtype=np.int64)
+        return verts
+
+    @property
+    def n_vertices(self) -> int:
+        return len(self.vertex_to_network)
+
+
+def load_yeo_atlas(cache_dir: str = "./cache/yeo") -> YeoAtlas:
+    """Load or generate Yeo atlas. Returns YeoAtlas instance."""
+    cache = Path(cache_dir)
+    cache.mkdir(parents=True, exist_ok=True)
+    label_json = cache / "yeo7_fsaverage5.json"
+
+    if label_json.exists():
+        return YeoAtlas(str(label_json))
+
+    # If user hasn't provided labels, create fallback and save it
+    atlas = YeoAtlas()
+    with open(label_json, "w") as f:
+        json.dump({str(k): int(v) for k, v in atlas.vertex_to_network.items()}, f)
+    return atlas
+
+
+class TribeV2Wrapper:
+    """
+    Wraps Meta's TRIBE v2 model for video → brain inference.
+    Falls back to a synthetic brain generator if TRIBE is not installed.
+    """
+
+    def __init__(self, model_id: str = "facebook/tribev2", device: str = "auto"):
+        self.model_id = model_id
+        self.device = device
+        self._model = None
+        self._has_tribe = False
+        self._init_model()
+
+    def _init_model(self):
+        if not HAS_TORCH:
+            warnings.warn("PyTorch not installed — using synthetic brain fallback.")
+            return
+        try:
+            from tribev2 import TribeModel
+            self._model = TribeModel.from_pretrained(self.model_id)
+            self._has_tribe = True
+            print(f"TRIBE v2 loaded from {self.model_id}")
+        except Exception as e:
+            warnings.warn(f"TRIBE v2 unavailable ({e}) — using synthetic brain fallback.")
+
+    def predict_brain(self, video_path: str, tr: float = 1.5) -> np.ndarray:
+        """
+        Run video through TRIBE v2 (or synthetic fallback).
+        Returns: (n_timesteps, n_vertices) array of predicted cortical response.
+        TRIBE native TR is ~1.49 s (Courtois NeuroMod). We resample to `tr` if needed.
+        """
+        if self._has_tribe and self._model is not None:
+            return self._predict_tribe(video_path, tr)
+        return self._predict_synthetic(video_path, tr)
+
+    def _predict_tribe(self, video_path: str, tr: float) -> np.ndarray:
+        from tribev2 import TribeModel
+        df = self._model.get_events_dataframe(video_path=video_path)
+        preds, segments = self._model.predict(events=df)
+        # preds shape: (n_timesteps, n_vertices)
+        # TRIBE TR is ~1.49s; if user wants different TR, we can resample later
+        return np.asarray(preds)
+
+    def _predict_synthetic(self, video_path: str, tr: float) -> np.ndarray:
+        """
+        Deterministic synthetic brain response for development/testing.
+        Uses video duration + hash of filename to seed a realistic-looking
+        BOLD-like time series per vertex.
+        """
+        import cv2
+        cap = cv2.VideoCapture(video_path)
+        fps = cap.get(cv2.CAP_PROP_FPS) or 30.0
+        frame_count = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+        duration = frame_count / fps if fps > 0 else 10.0
+        cap.release()
+
+        n_timesteps = max(1, int(np.ceil(duration / tr)))
+        n_vertices = 20484  # fsaverage5
+
+        # Seed from video name for reproducibility
+        seed = int.from_bytes(Path(video_path).name.encode(), "little") % (2**31)
+        rng = np.random.default_rng(seed)
+
+        # Base signal: slow drift + fast fluctuation per network
+        t = np.arange(n_timesteps) * tr
+        signal = np.zeros((n_timesteps, n_vertices), dtype=np.float32)
+
+        for net_id, net_name in YEO7_NAMES.items():
+            if net_id == 0:
+                continue
+            verts = np.where(np.arange(n_vertices) % 7 == (net_id - 1))[0]
+            # Network-specific temporal dynamics
+            freq = 0.1 + 0.05 * net_id  # ~0.01–0.1 Hz BOLD-ish
+            amp = 0.5 + 0.2 * np.sin(net_id)
+            base = amp * np.sin(2 * np.pi * freq * t + net_id)
+            noise = 0.1 * rng.standard_normal(n_timesteps)
+            signal[:, verts] = (base + noise).reshape(-1, 1)
+
+        return signal
+
+    def network_means(self, brain: np.ndarray, atlas: YeoAtlas) -> Dict[str, np.ndarray]:
+        """
+        Collapse vertex-wise brain response into per-Yeo-network time series.
+        brain: (n_timesteps, n_vertices)
+        Returns: {network_name: (n_timesteps,) array}
+        """
+        out = {}
+        for net_id, net_name in YEO7_NAMES.items():
+            if net_id == 0:
+                continue
+            verts = atlas.mask_for_network(net_name)
+            if len(verts):
+                out[net_name] = brain[:, verts].mean(axis=1)
+            else:
+                out[net_name] = np.zeros(brain.shape[0])
+        return out