pytractoviz 0.2.14__py3-none-any.whl

pytractoviz/utils.py

@@ -0,0 +1,220 @@
+"""Utility functions for tractography visualization."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import numpy as np
+
+if TYPE_CHECKING:
+    from dipy.tracking.streamline import Streamlines
+    from fury import window
+
+# Standard anatomical view angles (elevation, azimuth, roll)
+ANATOMICAL_VIEW_ANGLES = {
+    "coronal": (-90.0, 0.0, 0.0),  # Front view
+    "axial": (0.0, 180.0, 0.0),  # Top-down view
+    "sagittal": (-90.0, 0.0, 90.0),  # Side view (right)
+}
+
+
+def calculate_centroid(streamlines: Streamlines) -> np.ndarray:
+    """Calculate the centroid of streamlines.
+
+    Parameters
+    ----------
+    streamlines : Streamlines
+        The streamlines to calculate centroid for.
+
+    Returns
+    -------
+    np.ndarray
+        The centroid coordinates (3D).
+    """
+    all_points = np.vstack([np.array(sl) for sl in streamlines])
+    return np.mean(all_points, axis=0)
+
+
+def calculate_bbox_size(streamlines: Streamlines) -> np.ndarray:
+    """Calculate bounding box size of streamlines.
+
+    Parameters
+    ----------
+    streamlines : Streamlines
+        The streamlines to calculate bbox for.
+
+    Returns
+    -------
+    np.ndarray
+        The bounding box size (3D).
+    """
+    all_points = np.vstack([np.array(sl) for sl in streamlines])
+    return np.max(all_points, axis=0) - np.min(all_points, axis=0)
+
+
+def calculate_direction_colors(streamlines: Streamlines) -> np.ndarray:
+    """Calculate colors for streamlines based on their diffusion direction.
+
+    Standard RGB mapping:
+    - Red = X-axis (left/right)
+    - Green = Y-axis (anterior/posterior)
+    - Blue = Z-axis (superior/inferior)
+
+    Parameters
+    ----------
+    streamlines : Streamlines
+        The streamlines to calculate colors for.
+
+    Returns
+    -------
+    np.ndarray
+        Array of RGB colors, one per streamline (N x 3).
+    """
+    streamline_colors = []
+    max_range = 1e-10
+    sl_len = 2
+    for sl in streamlines:
+        sl_array = np.array(sl)
+        if len(sl_array) < sl_len:
+            # Degenerate streamline, use default color
+            streamline_colors.append([0.5, 0.5, 0.5])
+            continue
+
+        # Calculate direction vector (from start to end)
+        direction = sl_array[-1] - sl_array[0]
+        direction_norm = np.linalg.norm(direction)
+
+        if direction_norm < max_range:
+            # Degenerate direction, use default color
+            streamline_colors.append([0.5, 0.5, 0.5])
+            continue
+
+        # Normalize direction to unit vector
+        direction = direction / direction_norm
+
+        # Map direction components to RGB using absolute values
+        # X -> Red (left/right), Y -> Green (anterior/posterior), Z -> Blue (superior/inferior)
+        r = abs(direction[0])  # X component -> Red
+        g = abs(direction[1])  # Y component -> Green
+        b = abs(direction[2])  # Z component -> Blue
+
+        # Normalize by the maximum component to ensure colors are in 0-1 range
+        # This preserves the relative direction while ensuring valid RGB values
+        max_component = max(r, g, b)
+        if max_component > max_range:
+            r = r / max_component
+            g = g / max_component
+            b = b / max_component
+        else:
+            # Fallback for edge case
+            r, g, b = 0.5, 0.5, 0.5
+
+        streamline_colors.append([r, g, b])
+
+    return np.array(streamline_colors)
+
+
+def calculate_combined_centroid(*streamlines_groups: Streamlines) -> np.ndarray:
+    """Calculate the centroid of multiple groups of streamlines combined.
+
+    Parameters
+    ----------
+    *streamlines_groups : Streamlines
+        One or more Streamlines objects to combine.
+
+    Returns
+    -------
+    np.ndarray
+        The combined centroid coordinates (3D).
+    """
+    all_points_list = []
+    for streamlines in streamlines_groups:
+        all_points_list.append(np.vstack([np.array(sl) for sl in streamlines]))
+    all_points = np.vstack(all_points_list)
+    return np.mean(all_points, axis=0)
+
+
+def calculate_combined_bbox_size(*streamlines_groups: Streamlines) -> np.ndarray:
+    """Calculate bounding box size of multiple groups of streamlines combined.
+
+    Parameters
+    ----------
+    *streamlines_groups : Streamlines
+        One or more Streamlines objects to combine.
+
+    Returns
+    -------
+    np.ndarray
+        The combined bounding box size (3D).
+    """
+    all_points_list = []
+    for streamlines in streamlines_groups:
+        all_points_list.append(np.vstack([np.array(sl) for sl in streamlines]))
+    all_points = np.vstack(all_points_list)
+    return np.max(all_points, axis=0) - np.min(all_points, axis=0)
+
+
+def set_anatomical_camera(
+    scene: window.Scene,
+    centroid: np.ndarray,
+    view_name: str,
+    *,
+    camera_distance: float | None = None,
+    bbox_size: np.ndarray | None = None,
+) -> None:
+    """Set camera position for standard anatomical views.
+
+    This function positions the camera for coronal, axial, or sagittal views
+    without rotating the streamlines, ensuring colors stay aligned.
+
+    Parameters
+    ----------
+    scene : window.Scene
+        The FURY scene to set the camera on.
+    centroid : np.ndarray
+        The centroid of the streamlines (3D coordinates).
+    view_name : str
+        Name of the view: "coronal", "axial", or "sagittal".
+    camera_distance : float | None, optional
+        Distance of camera from centroid. If None, calculated from bbox_size.
+    bbox_size : np.ndarray | None, optional
+        Bounding box size of streamlines. Used to calculate camera_distance if not provided.
+
+    Raises
+    ------
+    ValueError
+        If view_name is not one of the standard anatomical views.
+    """
+    # Calculate camera distance if not provided
+    if camera_distance is None:
+        max_dim = np.max(bbox_size) if bbox_size is not None else 100.0
+        camera_distance = max_dim * 2.5
+
+    # Define camera positions and view_up vectors for each anatomical view
+    if view_name == "coronal":
+        # Coronal: front view (anterior), looking posterior
+        # Camera from posterior (-Y) to flip left/right, looking at centroid
+        camera_position = centroid + np.array([0, -camera_distance, 0])
+        view_up = np.array([0, 0, 1])  # Superior is up
+    elif view_name == "axial":
+        # Axial: top-down view (superior), looking inferior
+        # Camera from superior (+Z), looking at centroid
+        # Use -Y as view_up to flip left/right (instead of +Y)
+        camera_position = centroid + np.array([0, 0, camera_distance])
+        view_up = np.array([0, -1, 0])  # Posterior as up (flips left/right)
+    elif view_name == "sagittal":
+        # Sagittal: side view (left), looking right
+        # Camera from left (-X), looking at centroid
+        camera_position = centroid + np.array([-camera_distance, 0, 0])
+        view_up = np.array([0, 0, 1])  # Superior is up
+    else:
+        raise ValueError(
+            f"Invalid view name: {view_name}. Must be one of: coronal, axial, sagittal",
+        )
+
+    # Set camera
+    scene.set_camera(
+        position=camera_position,
+        focal_point=centroid,
+        view_up=view_up,
+    )