New upsampling method (PTT) for tractogram

scilpy/tractograms/tractogram_operations.py

@@ -19,12 +19,14 @@
 from dipy.tracking.streamline import transform_streamlines
 from nibabel.streamlines.array_sequence import ArraySequence
 import numpy as np
+from numpy.polynomial.polynomial import Polynomial
 from scipy.ndimage import map_coordinates
 from scipy.spatial import cKDTree
 
 from scilpy.tractograms.streamline_operations import smooth_line_gaussian, \
-    smooth_line_spline
+    smooth_line_spline, resample_streamlines_step_size
 from scilpy.utils.streamlines import cut_invalid_streamlines
+from scilpy.utils.spatial_ops import parallel_transport_streamline
 
 MIN_NB_POINTS = 10
 KEY_INDEX = np.concatenate((range(5), range(-1, -6, -1)))
@@ -603,7 +605,8 @@ def transform_warp_sft(sft, linear_transfo, target, inverse=False,
     return new_sft
 
 
-def upsample_tractogram(sft, nb, point_wise_std=None, streamline_wise_std=None,
+def upsample_tractogram(sft, nb, point_wise_std=None,
+                        streamline_wise_std=None, keep_tube=True,
                         gaussian=None, spline=None, seed=None):
     """
     Generates new streamlines by either adding gaussian noise around
@@ -622,6 +625,8 @@ def upsample_tractogram(sft, nb, point_wise_std=None, streamline_wise_std=None,
     streamline_wise_std : float
         The standard deviation of the gaussian to use to generate
         streamline-wise noise on the streamlines.
+    keep_tube : bool
+        If True, simply move the streamlines along their tangent uniformely.
     gaussian: float
         The sigma used for smoothing streamlines.
     spline: (float, int)
@@ -635,27 +640,54 @@ def upsample_tractogram(sft, nb, point_wise_std=None, streamline_wise_std=None,
     new_sft : StatefulTractogram
         The upsampled tractogram.
     """
-    assert bool(point_wise_std) ^ bool(streamline_wise_std), \
-        'Can only add either point-wise or streamline-wise noise' + \
-        ', not both nor none.'
+    def linear(data):
+        """Scale data linearly to [-1, 1]."""
+        min_val = np.min(data)
+        max_val = np.max(data)
+        return 2 * (data - min_val) / (max_val - min_val) - 1
+
+    def sigmoid(data):
+        """Scale data using sigmoid and then to [-1, 1]."""
+        return 2 * (1 / (1 + np.exp(-data))) - 1
+
+    def tanh(data):
+        """Scale data using tanh to [-1, 1]."""
+        return np.tanh(data)
+
+    def exp(data):
+        """Scale data using exp and then to [-1, 1]."""
+        exp_data = np.exp(data)
+        scaled_data = exp_data / np.max(exp_data)
+        return 2 * scaled_data - 1
 
     rng = np.random.RandomState(seed)
 
-    # Get the number of streamlines to add
-    nb_new = nb - len(sft.streamlines)
-
     # Get the streamlines that will serve as a base for new ones
-    indices = rng.choice(
-        len(sft.streamlines), nb_new)
-    new_streamlines = sft.streamlines.copy()
-
+    resample_sft = resample_streamlines_step_size(sft, 1)
+    new_streamlines = []
+    indices = np.random.choice(len(resample_sft), nb,
+                               replace=True)
     # For all selected streamlines, add noise and smooth
-    for s in sft.streamlines[indices]:
-        if point_wise_std:
-            noise = rng.normal(scale=point_wise_std, size=s.shape)
-        else:  # streamline_wise_std
-            noise = rng.normal(scale=streamline_wise_std, size=s.shape[-1])
-        new_s = s + noise
+    for s in resample_sft.streamlines[indices]:
+        if len(s) < 3:
+            new_streamlines.append(s)
+
+        new_s = parallel_transport_streamline(s, 1,
+                                              streamline_wise_std)[0]
+
+        # Generate smooth noise_factor
+        noise = np.random.normal(loc=0, scale=point_wise_std,
+                                 size=len(s))
+        x = np.arange(len(noise))
+        poly_coeffs = np.polyfit(x, noise, 3)
+        polynomial = Polynomial(poly_coeffs[::-1])
+        noise_factor = polynomial(x)
+        # print(noise_factor)
+
+        vec = s - new_s
+        vec /= np.linalg.norm(vec, axis=0)
+        new_s += vec * np.expand_dims(noise_factor, axis=1)
+
         if gaussian:
             new_s = smooth_line_gaussian(new_s, gaussian)
         elif spline:

scilpy/utils/spatial_ops.py

@@ -0,0 +1,153 @@
+# -*- coding: utf-8 -*-
+
+import numpy as np
+
+
+def normalize_vector(v):
+    """Normalize a 3D vector."""
+    norm = np.linalg.norm(v)
+    return v / norm if norm != 0 else v
+
+
+def find_perpendicular_vectors(v):
+    """Find two perpendicular unit vectors for a given normalized 3D vector."""
+    if v[0] == 0 and v[1] == 0:
+        if v[2] == 0:
+            # v is a zero vector, can't find perpendicular vectors
+            return None, None
+        # v is along the z-axis, choose x and y axes as perpendicular vectors
+        return np.array([1, 0, 0]), np.array([0, 1, 0])
+
+    # General case, find one vector perpendicular to v and z-axis
+    u = np.cross(v, [0, 0, 1])
+    u = normalize_vector(u)
+
+    # Find another vector perpendicular to both v and u
+    w = np.cross(v, u)
+    w = normalize_vector(w)
+
+    return u, w
+
+
+def project_on_plane(vector, u, w):
+    """Project a vector onto the plane defined by vectors u and w."""
+    # Calculate projections onto u and w
+    proj_u = np.dot(vector, u) * u
+    proj_w = np.dot(vector, w) * w
+
+    # Combine projections
+    projection = proj_u + proj_w
+
+    # Normalize and scale to match the original vector's norm
+    norm_vector = np.linalg.norm(vector)
+    normalized_projection = normalize_vector(projection)
+    scaled_projection = normalized_projection * norm_vector
+
+    return scaled_projection
+
+
+def parallel_transport_streamline(streamline, nb_streamlines, radius):
+    """ Generate new streamlines by parallel transport of the input
+    streamline. See [0] and [1] for more details.
+
+    [0]: Hanson, A.J., & Ma, H. (1995). Parallel Transport Approach to Curve Framing. # noqa E501
+    [1]: TD Essentials: Parallel Transport. https://www.youtube.com/watch?v=5LedteSEgOE
+
+    Parameters
+    ----------
+    streamline: ndarray (N, 3)
+        The streamline to transport.
+    nb_streamlines: int
+        The number of streamlines to generate.
+
+    Returns
+    -------
+    new_streamlines: list of ndarray (N, 3)
+        The generated streamlines.
+    """
+
+    def r(vec, theta):
+        """ Rotation matrix around a 3D vector by an angle theta.
+        From https://stackoverflow.com/questions/6802577/rotation-of-3d-vector
+
+        TODO?: Put this somewhere else.
+
+        Parameters
+        ----------
+        vec: ndarray (3,)
+            The vector to rotate around.
+        theta: float
+            The angle of rotation in radians.
+
+        Returns
+        -------
+        rot: ndarray (3, 3)
+            The rotation matrix.
+        """
+
+        vec = vec / np.linalg.norm(vec)
+        x, y, z = vec
+        c, s = np.cos(theta), np.sin(theta)
+        return np.array([[c + x**2 * (1 - c),
+                          x * y * (1 - c) - z * s,
+                          x * z * (1 - c) + y * s],
+                         [y * x * (1 - c) + z * s,
+                          c + y**2 * (1 - c),
+                          y * z * (1 - c) - x * s],
+                         [z * x * (1 - c) - y * s,
+                          z * y * (1 - c) + x * s,
+                          c + z**2 * (1 - c)]])
+
+    # Compute the tangent at each point of the streamline
+    T = np.gradient(streamline, axis=0)
+    # Normalize the tangents
+    T = T / np.linalg.norm(T, axis=1)[:, None]
+
+    # Placeholder for the normal vector at each point
+    V = np.zeros_like(T)
+    # Set the normal vector at the first point to be [0, 1, 0]
+    # (arbitrary choice)
+    V[0] = np.random.rand(3)
+
+    # For each point
+    for i in range(0, T.shape[0]-1):
+        # Compute the torsion vector
+        B = np.cross(T[i], T[i+1])
+        # If the torsion vector is 0, the normal vector does not change
+        if np.linalg.norm(B) < 1e-3:
+            V[i+1] = V[i]
+        # Else, the normal vector is rotated around the torsion vector by
+        # the torsion.
+        else:
+            B = B / np.linalg.norm(B)
+            theta = np.arccos(np.dot(T[i], T[i+1]))
+            # Rotate the vector V[i] around the vector B by theta
+            # radians.
+            V[i+1] = np.dot(r(B, theta), V[i])
+
+    # Compute the binormal vector at each point
+    W = np.cross(T, V, axis=1)
+
+    # Generate the new streamlines
+    # TODO?: This could easily be optimized to avoid the for loop, we have to
+    # see if this becomes a bottleneck.
+    new_streamlines = []
+    for i in range(nb_streamlines):
+        # Get a random number between -1 and 1
+        rand_v = (np.random.rand() * 2 - 1)
+        rand_w = (np.random.rand() * 2 - 1)
+        # Compute the norm of the "displacement"
+        norm = np.sqrt(rand_v**2 + rand_w**2)
+        # Displace the normal and binormal vectors by a random amount
+        V_mod = V * rand_v
+        W_mod = W * rand_w
+        # Compute the displacement vector
+        VW = (V_mod + W_mod)
+        # Displace the streamline around the original one following the
+        # parallel frame. Make sure to normalize the displacement vector
+        # so that the new streamline is in a circle around the original one.
+
+        new_s = streamline + (np.random.rand() * VW / norm) * radius
+        new_streamlines.append(new_s)
+
+    return new_streamlines

scilpy/utils/util.py

@@ -160,4 +160,4 @@ def recursive_print(data):
         print(list(data.keys()))
         recursive_print(data[list(data.keys())[0]])
     else:
-        return
+        return

scripts/scil_resample_tractogram.py

@@ -61,13 +61,14 @@ def _build_arg_parser():
 
     # For upsampling:
     upsampling_group = p.add_argument_group('Upsampling params')
-    std_group = upsampling_group.add_mutually_exclusive_group()
-    std_group.add_argument('--point_wise_std', type=float,
+    upsampling_group.add_argument('--point_wise_std', type=float, default=1,
                            help='Noise to add to existing streamlines\'' +
-                                ' points to generate new ones.')
-    std_group.add_argument('--streamline_wise_std', type=float,
+                                ' points to generate new ones [%(default)s].')
+    upsampling_group.add_argument('--streamline_wise_std', type=float, default=1,
                            help='Noise to add to existing whole' +
-                                ' streamlines to generate new ones.')
+                                ' streamlines to generate new ones [%(default)s].')
+    upsampling_group.add_argument('--keep_tube', action='store_true',
+                            help='Keep streamlines as tube (default: False).')
     sub_p = upsampling_group.add_mutually_exclusive_group()
     sub_p.add_argument('--gaussian', metavar='SIGMA', type=int,
                        help='Sigma for smoothing. Use the value of surronding'
@@ -146,6 +147,7 @@ def main():
         sft = upsample_tractogram(
             sft, args.nb_streamlines,
             args.point_wise_std, args.streamline_wise_std,
+            args.keep_tube,
             args.gaussian, args.spline, args.seed)
     elif args.nb_streamlines < original_number:
         if args.downsample_per_cluster: