Upload live_engine.py with huggingface_hub

live_engine.py

@@ -0,0 +1,284 @@
+"""Real-time brain prediction engine.
+
+Runs in a background thread, consuming frames from a capture source,
+extracting features, and producing brain predictions via TRIBE v2.
+
+When CortexLab is not installed, falls back to a simulation mode that
+generates synthetic predictions from frame statistics.
+"""
+
+from __future__ import annotations
+
+import time
+import threading
+import logging
+from collections import deque
+from dataclasses import dataclass, field
+
+import numpy as np
+
+from live_capture import BaseCapture, MediaFrame
+
+logger = logging.getLogger(__name__)
+
+# Check if CortexLab is available
+try:
+    from cortexlab.inference.predictor import TribeModel
+    CORTEXLAB_AVAILABLE = True
+except ImportError:
+    CORTEXLAB_AVAILABLE = False
+
+
+@dataclass
+class LivePrediction:
+    """A single prediction with metadata."""
+    vertex_data: np.ndarray  # (n_vertices,)
+    timestamp: float
+    cognitive_load: dict[str, float] = field(default_factory=dict)
+    processing_time_ms: float = 0.0
+
+
+@dataclass
+class LiveMetrics:
+    """Aggregated metrics from the live engine."""
+    fps: float = 0.0
+    total_frames: int = 0
+    total_predictions: int = 0
+    avg_latency_ms: float = 0.0
+    is_running: bool = False
+    mode: str = "simulation"  # "simulation" or "cortexlab"
+
+
+class LiveInferenceEngine:
+    """Background engine for real-time brain prediction.
+
+    Consumes frames from a capture source and produces brain predictions.
+    If CortexLab is installed and a GPU is available, uses the real TRIBE v2
+    model. Otherwise, falls back to simulation mode that generates plausible
+    predictions from frame statistics.
+    """
+
+    def __init__(
+        self,
+        n_vertices: int = 580,
+        roi_indices: dict | None = None,
+        buffer_size: int = 120,
+        checkpoint: str = "facebook/tribev2",
+        device: str = "auto",
+        cache_folder: str = "./cache",
+    ):
+        self.n_vertices = n_vertices
+        self.roi_indices = roi_indices or {}
+        self.buffer_size = buffer_size
+        self.checkpoint = checkpoint
+        self.device = device
+        self.cache_folder = cache_folder
+
+        self._predictions: deque[LivePrediction] = deque(maxlen=buffer_size)
+        self._running = False
+        self._thread: threading.Thread | None = None
+        self._lock = threading.Lock()
+        self._model = None
+        self._metrics = LiveMetrics()
+        self._capture: BaseCapture | None = None
+
+    def start(self, capture: BaseCapture):
+        """Start the inference engine with a media capture source."""
+        if self._running:
+            return
+
+        self._capture = capture
+        self._running = True
+        self._metrics = LiveMetrics(is_running=True)
+
+        # Try to load CortexLab model
+        if CORTEXLAB_AVAILABLE:
+            try:
+                logger.info("Loading TRIBE v2 model...")
+                self._model = TribeModel.from_pretrained(
+                    self.checkpoint, device=self.device, cache_folder=self.cache_folder
+                )
+                self._metrics.mode = "cortexlab"
+                logger.info("Model loaded. Using real inference.")
+            except Exception as e:
+                logger.warning(f"Failed to load model: {e}. Using simulation mode.")
+                self._model = None
+                self._metrics.mode = "simulation"
+        else:
+            self._metrics.mode = "simulation"
+
+        capture.start()
+        self._thread = threading.Thread(target=self._inference_loop, daemon=True)
+        self._thread.start()
+
+    def stop(self):
+        """Stop the engine and capture source."""
+        self._running = False
+        if self._capture:
+            self._capture.stop()
+        if self._thread:
+            self._thread.join(timeout=5.0)
+        self._metrics.is_running = False
+
+    def get_latest_prediction(self) -> LivePrediction | None:
+        with self._lock:
+            return self._predictions[-1] if self._predictions else None
+
+    def get_predictions(self, n: int = 60) -> list[LivePrediction]:
+        with self._lock:
+            return list(self._predictions)[-n:]
+
+    def get_metrics(self) -> LiveMetrics:
+        return self._metrics
+
+    def _inference_loop(self):
+        """Main loop: consume frames, produce predictions."""
+        frame_times = deque(maxlen=30)
+        last_frame_count = 0
+
+        while self._running:
+            frame = self._capture.get_latest_frame()
+            if frame is None:
+                time.sleep(0.1)
+                continue
+
+            # Skip if we already processed this frame
+            current_count = self._capture.frame_count
+            if current_count == last_frame_count:
+                time.sleep(0.05)
+                continue
+            last_frame_count = current_count
+
+            start = time.time()
+
+            if self._model is not None and self._metrics.mode == "cortexlab":
+                prediction = self._run_real_inference(frame)
+            else:
+                prediction = self._run_simulation(frame)
+
+            elapsed_ms = (time.time() - start) * 1000
+            prediction.processing_time_ms = elapsed_ms
+
+            with self._lock:
+                self._predictions.append(prediction)
+
+            # Update metrics
+            frame_times.append(time.time())
+            self._metrics.total_predictions += 1
+            self._metrics.total_frames = current_count
+            self._metrics.avg_latency_ms = elapsed_ms
+            if len(frame_times) >= 2:
+                self._metrics.fps = (len(frame_times) - 1) / (frame_times[-1] - frame_times[0])
+
+        # Check if capture stopped (file ended)
+        if not self._capture.is_running:
+            self._running = False
+            self._metrics.is_running = False
+
+    def _run_real_inference(self, frame: MediaFrame) -> LivePrediction:
+        """Run actual TRIBE v2 inference on a frame.
+
+        For real-time, we skip the full pipeline (get_events_dataframe)
+        and use a simplified feature extraction path.
+        """
+        import tempfile
+        import os
+
+        try:
+            # Save frame as temporary video (1 frame)
+            import cv2
+            tmp_path = os.path.join(tempfile.gettempdir(), "cortexlab_live_frame.mp4")
+            h, w = frame.video_frame.shape[:2]
+            out = cv2.VideoWriter(tmp_path, cv2.VideoWriter_fourcc(*'mp4v'), 1, (w, h))
+            out.write(cv2.cvtColor(frame.video_frame, cv2.COLOR_RGB2BGR))
+            out.release()
+
+            events = self._model.get_events_dataframe(video_path=tmp_path)
+            preds, _ = self._model.predict(events, verbose=False)
+            vertex_data = preds.mean(axis=0) if preds.ndim == 2 else preds
+
+            # Normalize to [0, 1]
+            vmin, vmax = vertex_data.min(), vertex_data.max()
+            if vmax > vmin:
+                vertex_data = (vertex_data - vmin) / (vmax - vmin)
+
+            os.unlink(tmp_path)
+        except Exception as e:
+            logger.warning(f"Inference failed: {e}. Falling back to simulation.")
+            return self._run_simulation(frame)
+
+        cog_load = self._compute_cognitive_load(vertex_data)
+        return LivePrediction(
+            vertex_data=vertex_data,
+            timestamp=frame.timestamp,
+            cognitive_load=cog_load,
+        )
+
+    def _run_simulation(self, frame: MediaFrame) -> LivePrediction:
+        """Generate plausible predictions from frame statistics.
+
+        Uses frame brightness/color as proxy for visual complexity,
+        creating biologically-inspired activation patterns.
+        """
+        rng = np.random.default_rng(int(frame.timestamp * 1000) % (2**31))
+
+        # Base noise
+        vertex_data = rng.standard_normal(self.n_vertices) * 0.03
+
+        if frame.video_frame is not None:
+            img = frame.video_frame.astype(np.float32) / 255.0
+
+            # Visual complexity from image statistics
+            brightness = img.mean()
+            contrast = img.std()
+            color_variance = img.var(axis=(0, 1)).mean()
+
+            # Map to ROI activations
+            for roi_name, vertices in self.roi_indices.items():
+                valid = vertices[vertices < self.n_vertices]
+                if len(valid) == 0:
+                    continue
+
+                # Visual ROIs respond to brightness/contrast
+                if roi_name in ["V1", "V2", "V3", "V4", "MT", "MST", "FFC", "VVC"]:
+                    activation = contrast * 0.8 + color_variance * 0.5
+                # Auditory ROIs get low baseline
+                elif roi_name in ["A1", "LBelt", "MBelt", "PBelt", "A4", "A5"]:
+                    activation = 0.05 + rng.random() * 0.1
+                # Language ROIs moderate
+                elif roi_name in ["44", "45", "IFJa", "IFJp", "TPOJ1", "TPOJ2"]:
+                    activation = brightness * 0.3
+                # Executive ROIs track change
+                elif roi_name in ["46", "9-46d", "8Av", "8Ad", "FEF"]:
+                    activation = contrast * 0.5
+                else:
+                    activation = 0.1
+
+                vertex_data[valid] = activation + rng.standard_normal(len(valid)) * 0.05
+
+        vertex_data = np.clip(vertex_data, 0, 1)
+        cog_load = self._compute_cognitive_load(vertex_data)
+
+        return LivePrediction(
+            vertex_data=vertex_data,
+            timestamp=frame.timestamp,
+            cognitive_load=cog_load,
+        )
+
+    def _compute_cognitive_load(self, vertex_data: np.ndarray) -> dict[str, float]:
+        """Compute cognitive load dimensions from vertex data."""
+        from utils import COGNITIVE_DIMENSIONS
+
+        baseline = max(float(np.median(np.abs(vertex_data))), 1e-8)
+        scores = {}
+        for dim, rois in COGNITIVE_DIMENSIONS.items():
+            vals = []
+            for roi in rois:
+                if roi in self.roi_indices:
+                    verts = self.roi_indices[roi]
+                    valid = verts[verts < len(vertex_data)]
+                    if len(valid) > 0:
+                        vals.append(np.abs(vertex_data[valid]).mean())
+            scores[dim] = min(float(np.mean(vals)) / baseline, 1.0) if vals else 0.0
+        scores["Overall"] = float(np.mean(list(scores.values()))) if scores else 0.0
+        return scores