Merge pull request #35 from 36000/minmaxlen

[ENH] add min/max len

Author: 36000

cuslines/cuda_python/cu_propagate_seeds.py

@@ -25,9 +25,11 @@
 
 
 class SeedBatchPropagator:
-    def __init__(self, gpu_tracker):
+    def __init__(self, gpu_tracker, minlen=0, maxlen=np.inf):
         self.gpu_tracker = gpu_tracker
         self.ngpus = gpu_tracker.ngpus
+        self.minlen = minlen
+        self.maxlen = maxlen
 
         self.nSlines_old = np.zeros(self.ngpus, dtype=np.int32)
         self.nSlines = np.zeros(self.ngpus, dtype=np.int32)
@@ -240,6 +242,8 @@ def get_buffer_size(self):
         for ii in range(self.ngpus):
             lens = self.sline_lens[ii]
             for jj in range(self.nSlines[ii]):
+                if lens[jj] < self.minlen or lens[jj] > self.maxlen:
+                    continue
                 buffer_size += lens[jj] * 3 * REAL_SIZE
         return math.ceil(buffer_size / MEGABYTE)
 
@@ -252,9 +256,14 @@ def _yield_slines():
                 for jj in range(self.nSlines[ii]):
                     npts = this_len[jj]
 
+                    if npts < self.minlen or npts > self.maxlen:
+                        continue
+
                     yield np.asarray(this_sls[jj], dtype=REAL_DTYPE)[:npts]
 
         return _yield_slines()
 
     def as_array_sequence(self):
-        return ArraySequence(self.as_generator(), self.get_buffer_size())
+        return ArraySequence(
+            self.as_generator(),
+            self.get_buffer_size())

cuslines/cuda_python/cu_tractography.py

@@ -44,6 +44,8 @@ def __init__(
         sphere_edges: np.ndarray,
         max_angle: float = radians(60),
         step_size: float = 0.5,
+        min_pts=0,
+        max_pts=np.inf,
         relative_peak_thresh: float = 0.25,
         min_separation_angle: float = radians(45),
         ngpus: int = 1,
@@ -74,8 +76,14 @@ def __init__(
             Maximum angle (in radians) between steps
             default: radians(60)
         step_size : float, optional
-            Step size for tracking
+            Step size for tracking, in voxels
             default: 0.5
+        min_pts : int, optional
+            Minimum number of points in a streamline to be kept
+            default: 0
+        max_pts : int, optional
+            Maximum number of points in a streamline to be kept
+            default: np.inf
         relative_peak_thresh : float, optional
             Relative peak threshold for direction selection
             default: 0.25
@@ -136,7 +144,11 @@ def __init__(
         self.streams = []
         self.managed_data = []
 
-        self.seed_propagator = SeedBatchPropagator(gpu_tracker=self)
+        self.seed_propagator = SeedBatchPropagator(
+            gpu_tracker=self,
+            minlen=min_pts,
+            maxlen=max_pts
+        )
         self._allocated = False
 
     def __enter__(self):
@@ -260,7 +272,7 @@ def generate_trx(self, seeds, ref_img):
 
         # Will resize by a factor of 2 if these are exceeded
         sl_len_guess = 100
-        sl_per_seed_guess = 4
+        sl_per_seed_guess = 2
         n_sls_guess = sl_per_seed_guess * seeds.shape[0]
 
         # trx files use memory mapping