def test_compress_streamlines_memory_leaks():
# Test some dtypes
dtypes = [np.float32, np.float64, np.int32, np.int64]
for dtype in dtypes:
rng = np.random.RandomState(1234)
NB_STREAMLINES = 10000
streamlines = [rng.randn(rng.randint(10, 100), 3).astype(dtype)
for _ in range(NB_STREAMLINES)]
list_refcount_before = get_type_refcount()["list"]
cstreamlines = compress_streamlines(streamlines)
list_refcount_after = get_type_refcount()["list"]
del cstreamlines # Delete `cstreamlines` because it holds a reference
# to `list`.
# Calling `compress_streamlines` should increase the refcount of `list`
# by one since we kept the returned value.
assert_equal(list_refcount_after, list_refcount_before+1)
# Test mixed dtypes
rng = np.random.RandomState(1234)
NB_STREAMLINES = 10000
streamlines = []
for i in range(NB_STREAMLINES):
dtype = dtypes[i % len(dtypes)]
streamlines.append(rng.randn(rng.randint(10, 100), 3).astype(dtype))
list_refcount_before = get_type_refcount()["list"]
cstreamlines = compress_streamlines(streamlines)
list_refcount_after = get_type_refcount()["list"]
# Calling `compress_streamlines` should increase the refcount of `list` by
# one since we kept the returned value.
assert_equal(list_refcount_after, list_refcount_before+1)
def test_compress_streamlines_memory_leaks():
# Test some dtypes
dtypes = [np.float32, np.float64, np.int32, np.int64]
for dtype in dtypes:
rng = np.random.RandomState(1234)
NB_STREAMLINES = 10000
streamlines = [rng.randn(rng.randint(10, 100), 3).astype(dtype) for _ in range(NB_STREAMLINES)]
list_refcount_before = get_type_refcount()["list"]
cstreamlines = compress_streamlines(streamlines)
list_refcount_after = get_type_refcount()["list"]
del cstreamlines # Delete `cstreamlines` because it holds a reference to `list`.
# Calling `compress_streamlines` should increase the refcount of `list` by one
# since we kept the returned value.
assert_equal(list_refcount_after, list_refcount_before+1)
# Test mixed dtypes
rng = np.random.RandomState(1234)
NB_STREAMLINES = 10000
streamlines = []
for i in range(NB_STREAMLINES):
dtype = dtypes[i % len(dtypes)]
streamlines.append(rng.randn(rng.randint(10, 100), 3).astype(dtype))
list_refcount_before = get_type_refcount()["list"]
cstreamlines = compress_streamlines(streamlines)
list_refcount_after = get_type_refcount()["list"]
# Calling `compress_streamlines` should increase the refcount of `list` by one
# since we kept the returned value.
assert_equal(list_refcount_after, list_refcount_before+1)
def compress_fibers_worker_shared_mem(idx):
# Function that runs in parallel must be on top level (not in class/function) otherwise it can not be pickled and then error
streamlines_chunk = FIBER_BATCHES[
idx] # shared memory; by using indices each worker accesses only his part
result = compress_streamlines(streamlines_chunk,
tol_error=COMPRESSION_ERROR_THRESHOLD)
logging.debug('PID {}, DONE'.format(getpid()))
return result
def harvest(
self,
states: np.ndarray,
compress=False,
) -> Tuple[StatefulTractogram, np.ndarray]:
"""Internally keep only the streamlines and corresponding env. states
that haven't stopped yet, and return the streamlines that triggered a
stopping flag.
Parameters
----------
states: torch.Tensor
Environment states to be "pruned"
Returns
-------
tractogram : nib.streamlines.Tractogram
Tractogram containing the streamlines that stopped tracking,
along with the stopping_flags information and seeds in
`tractogram.data_per_streamline`
states: np.ndarray of size [n_streamlines, input_size]
Input size for all continuing last streamline positions and
neighbors + input addons
stopping_idx: np.ndarray
Indexes of stopping trajectories. Returned in case an RL
algorithm would need 'em
"""
tractogram = Tractogram()
# Harvest stopped streamlines and associated data
stopped_seeds = self.starting_points[self.stopping_idx]
stopped_streamlines = self.streamlines[self.stopping_idx, :self.length]
# Drop last point if it triggered a flag we don't want
flags = is_flag_set(self.stopping_flags,
StoppingFlags.STOPPING_CURVATURE)
streamlines = [
s[:-1] if f else s for f, s in zip(flags, stopped_streamlines)
]
if compress:
streamlines = compress_streamlines(streamlines, 0.1)
# Harvested tractogram
tractogram = Tractogram(streamlines=streamlines,
data_per_streamline={
"stopping_flags": self.stopping_flags,
"seeds": stopped_seeds
},
affine_to_rasmm=self.affine_vox2rasmm)
# Keep only streamlines that should continue
states = self._keep(self.continue_idx, states)
return tractogram, states, self.continue_idx
def extend(self, roi_sl):
# compress the new streamlines to save space
for i in range(len(roi_sl.streamlines)):
streamline = roi_sl.streamlines[i]
roi_sl.streamlines[i] = compress_streamlines(
streamline, COMPRESSION_RATE)
# add streamlines and intersection information
self.streamlines.extend(roi_sl.streamlines)
self.ids_in.extend(roi_sl.ids_in)
self.ids_out.extend(roi_sl.ids_out)
self.pts_in.extend(roi_sl.pts_in)
self.pts_out.extend(roi_sl.pts_out)
self.surf_in.extend(roi_sl.surf_in)
self.surf_out.extend(roi_sl.surf_out)
def harvest(self):
undone, done, stopping_flags = self.is_stopping(self.sprouts)
# Do not keep last point since it almost surely raised the stopping flag.
streamlines = list(self.sprouts[done, :-1])
if self.compress_streamlines:
streamlines = compress_streamlines(streamlines)
tractogram = Tractogram(streamlines=streamlines,
data_per_streamline={"stopping_flags": stopping_flags})
# Keep only undone sprouts
self._keep(undone)
return tractogram
def harvest(self):
undone, done, stopping_flags = self.is_stopping(self.sprouts)
tractogram = Tractogram(streamlines=compress_streamlines(list(self.sprouts[done])),
data_per_streamline={"stopping_flags": stopping_flags})
# Update remaining sprouts and their states.
self.sprouts = self.sprouts[undone]
self._states = [s[undone] for s in self._states]
# Rewrite history
for i, state in enumerate(self._history):
self._history[i] = [s[undone] for s in state]
return tractogram
def harvest(self):
undone, done, stopping_flags = self.is_stopping(
self.sprouts, self.sprouts_stop)
# Do not keep last point since it almost surely raised the stopping flag.
streamlines = list(self.sprouts[done, :-1])
if self.compress_streamlines:
streamlines = compress_streamlines(streamlines)
tractogram = Tractogram(
streamlines=streamlines,
data_per_streamline={"stopping_flags": stopping_flags})
# Keep only undone sprouts
self._keep(undone)
return tractogram
def dpy(workingdir,
input,
compressed=None,
restored=None,
tol_error=0.01,
force=False,
coords_only=False,
verbose=False):
if not input.endswith('tck') and not input.endswith('trk'):
# we need to convert
print('Invalid format')
return None
if not compressed:
compressed = input + '_compressed' + str(tol_error) + '.dpy'
if not restored:
restored = input + '_restored' + str(tol_error) + '.tck'
original = os.path.join(workingdir, input)
#
# compression
#
c_time = -1
if not os.path.exists(os.path.join(workingdir, compressed)) or force:
# compress again!
t0 = time.time()
loaded_original = nib.streamlines.load(original, lazy_load=False)
original_streamlines = loaded_original.streamlines
# parameters from Presseau15 are 0.01 and inf
c_streamlines = compress_streamlines(original_streamlines,
tol_error=tol_error,
max_segment_length=np.inf)
# write dpy file
# set compression to highest but it does have any effect
dpw = Dpy(os.path.join(workingdir, compressed),
mode='w',
compression=9)
for c_s in c_streamlines:
dpw.write_track(c_s)
dpw.close()
c_time = time.time() - t0
if verbose:
print('compression done.')
#
# restoring
#
d_time = -1
if not os.path.exists(os.path.join(workingdir, restored)) or force:
# restore again!
t0 = time.time()
restored_data = Dpy(os.path.join(workingdir, compressed), mode='r')
restored_streamlines = restored_data.read_tracks()
restored_data.close()
d_time = time.time() - t0
with open(os.path.join(workingdir, restored), 'w') as f:
f.write('restoredok.')
if verbose:
print('restoring done.')
#
# calculate errors
#
stats = compressed + '_stats' + str(tol_error) + '.p'
if not os.path.exists(os.path.join(workingdir, stats)) or force:
statsdata = Runner.error_per_streamlines(original_streamlines,
restored_streamlines)
sizestatsdata = Runner.sizestats(
os.path.join(workingdir, original),
os.path.join(workingdir, compressed))
statsdata = [
c_time, d_time, sizestatsdata, statsdata[0], statsdata[1]
]
with open(os.path.join(workingdir, stats), 'wb') as f:
pickle.dump(statsdata, f)
else:
with open(os.path.join(workingdir, stats), 'rb') as f:
statsdata = pickle.load(f)
[c_time, d_time, sizestatsdata, (min_e, max_e, mean_e, std_e), (end_min_e, end_max_e, end_mean_e, end_std_e)] = \
statsdata
if verbose:
print('Times', c_time, d_time)
print('Size Stats', sizestatsdata)
print('Error', min_e, max_e, mean_e, std_e)
print('Endpoint Error', end_min_e, end_max_e, end_mean_e,
end_std_e)
return statsdata
def test_compress_streamlines():
for compress_func in [compress_streamlines_python, compress_streamlines]:
# Small streamlines (less than two points) are uncompressable.
for small_streamline in [np.array([[]]),
np.array([[1, 1, 1]]),
np.array([[1, 1, 1], [2, 2, 2]])]:
c_streamline = compress_func(small_streamline)
assert_equal(len(c_streamline), len(small_streamline))
assert_array_equal(c_streamline, small_streamline)
# Compressing a straight streamline that is less than 10mm long
# should output a two points streamline.
linear_streamline = np.linspace(0, 5, 100*3).reshape((100, 3))
c_streamline = compress_func(linear_streamline)
assert_equal(len(c_streamline), 2)
assert_array_equal(c_streamline, [linear_streamline[0],
linear_streamline[-1]])
# The distance of consecutive points must be less or equal than some
# value.
max_segment_length = 10
linear_streamline = np.linspace(0, 100, 100*3).reshape((100, 3))
linear_streamline[:, 1:] = 0.
c_streamline = compress_func(linear_streamline,
max_segment_length=max_segment_length)
segments_length = np.sqrt((np.diff(c_streamline,
axis=0)**2).sum(axis=1))
assert_true(np.all(segments_length <= max_segment_length))
assert_equal(len(c_streamline), 12)
assert_array_equal(c_streamline, linear_streamline[::9])
# A small `max_segment_length` should keep all points.
c_streamline = compress_func(linear_streamline,
max_segment_length=0.01)
assert_array_equal(c_streamline, linear_streamline)
# Test we can set `max_segment_length` to infinity
# (like the C++ version)
compress_func(streamline, max_segment_length=np.inf)
# Uncompressable streamline when `tol_error` == 1.
simple_streamline = np.array([[0, 0, 0],
[1, 1, 0],
[1.5, np.inf, 0],
[2, 2, 0],
[2.5, 20, 0],
[3, 3, 0]])
# Because of np.inf, compressing that streamline causes a warning.
with np.errstate(invalid='ignore'):
c_streamline = compress_func(simple_streamline, tol_error=1)
assert_array_equal(c_streamline, simple_streamline)
# Create a special streamline where every other point is increasingly
# farther from a straigth line formed by the streamline endpoints.
tol_errors = np.linspace(0, 10, 21)
orthogonal_line = np.array([[-np.sqrt(2)/2, np.sqrt(2)/2, 0]],
dtype=np.float32)
special_streamline = np.array([range(len(tol_errors)*2+1)] * 3,
dtype=np.float32).T
special_streamline[1::2] += orthogonal_line * tol_errors[:, None]
# # Uncomment to see the streamline.
# import pylab as plt
# plt.plot(special_streamline[:, 0], special_streamline[:, 1], '.-')
# plt.axis('equal'); plt.show()
# Test different values for `tol_error`.
for i, tol_error in enumerate(tol_errors):
cspecial_streamline = compress_streamlines(special_streamline,
tol_error=tol_error+1e-4,
max_segment_length=np.inf)
# First and last points should always be the same as the original ones.
assert_array_equal(cspecial_streamline[0], special_streamline[0])
assert_array_equal(cspecial_streamline[-1], special_streamline[-1])
assert_equal(len(cspecial_streamline),
len(special_streamline)-((i*2)+1))
# Make sure Cython and Python versions are the same.
cstreamline_python = compress_streamlines_python(
special_streamline,
tol_error=tol_error+1e-4,
max_segment_length=np.inf)
assert_equal(len(cspecial_streamline), len(cstreamline_python))
assert_array_almost_equal(cspecial_streamline, cstreamline_python)
def test_compress_streamlines():
for compress_func in [compress_streamlines_python, compress_streamlines]:
# Small streamlines (less than two points) are uncompressable.
for small_streamline in [
np.array([[]]),
np.array([[1, 1, 1]]),
np.array([[1, 1, 1], [2, 2, 2]])
]:
c_streamline = compress_func(small_streamline)
assert_equal(len(c_streamline), len(small_streamline))
assert_array_equal(c_streamline, small_streamline)
# Compressing a straight streamline that is less than 10mm long
# should output a two points streamline.
linear_streamline = np.linspace(0, 5, 100 * 3).reshape((100, 3))
c_streamline = compress_func(linear_streamline)
assert_equal(len(c_streamline), 2)
assert_array_equal(c_streamline,
[linear_streamline[0], linear_streamline[-1]])
# The distance of consecutive points must be less or equal than some
# value.
max_segment_length = 10
linear_streamline = np.linspace(0, 100, 100 * 3).reshape((100, 3))
linear_streamline[:, 1:] = 0.
c_streamline = compress_func(linear_streamline,
max_segment_length=max_segment_length)
segments_length = np.sqrt((np.diff(c_streamline,
axis=0)**2).sum(axis=1))
assert_true(np.all(segments_length <= max_segment_length))
assert_equal(len(c_streamline), 12)
assert_array_equal(c_streamline, linear_streamline[::9])
# A small `max_segment_length` should keep all points.
c_streamline = compress_func(linear_streamline,
max_segment_length=0.01)
assert_array_equal(c_streamline, linear_streamline)
# Test we can set `max_segment_length` to infinity
# (like the C++ version)
compress_func(streamline, max_segment_length=np.inf)
# Uncompressable streamline when `tol_error` == 1.
simple_streamline = np.array([[0, 0, 0], [1, 1, 0], [1.5, np.inf, 0],
[2, 2, 0], [2.5, 20, 0], [3, 3, 0]])
# Because of np.inf, compressing that streamline causes a warning.
with np.errstate(invalid='ignore'):
c_streamline = compress_func(simple_streamline, tol_error=1)
assert_array_equal(c_streamline, simple_streamline)
# Create a special streamline where every other point is increasingly
# farther from a straigth line formed by the streamline endpoints.
tol_errors = np.linspace(0, 10, 21)
orthogonal_line = np.array(
[[-np.sqrt(2) / 2, np.sqrt(2) / 2, 0]], dtype=np.float32)
special_streamline = np.array([range(len(tol_errors) * 2 + 1)] * 3,
dtype=np.float32).T
special_streamline[1::2] += orthogonal_line * tol_errors[:, None]
# # Uncomment to see the streamline.
# import pylab as plt
# plt.plot(special_streamline[:, 0], special_streamline[:, 1], '.-')
# plt.axis('equal'); plt.show()
# Test different values for `tol_error`.
for i, tol_error in enumerate(tol_errors):
cspecial_streamline = compress_streamlines(special_streamline,
tol_error=tol_error + 1e-4,
max_segment_length=np.inf)
# First and last points should always be the same as the original ones.
assert_array_equal(cspecial_streamline[0], special_streamline[0])
assert_array_equal(cspecial_streamline[-1], special_streamline[-1])
assert_equal(len(cspecial_streamline),
len(special_streamline) - ((i * 2) + 1))
# Make sure Cython and Python versions are the same.
cstreamline_python = compress_streamlines_python(
special_streamline,
tol_error=tol_error + 1e-4,
max_segment_length=np.inf)
assert_equal(len(cspecial_streamline), len(cstreamline_python))
assert_array_almost_equal(cspecial_streamline, cstreamline_python)
def compress_fibers_worker_shared_mem(idx):
# Function that runs in parallel must be on top level (not in class/function) otherwise it can not be pickled and then error
streamlines_chunk = FIBER_BATCHES[idx] # shared memory; by using indices each worker accesses only his part
result = compress_streamlines(streamlines_chunk, tol_error=COMPRESSION_ERROR_THRESHOLD)
logging.debug('PID {}, DONE'.format(getpid()))
return result
mesh = TriMesh_Vtk(args.surface, None)
end_vertices = mesh.positive_mass_stiffness_smooth(
nb_iter=args.nb_step,
diffusion_step=step_size,
flow_file=args.flow_file)
end_mesh = TriMesh_Vtk(mesh.get_triangles(), end_vertices)
print "saving surface tracking ..."
flow = np.memmap(args.flow_file, dtype=np.float64, mode='r', shape=(args.nb_step, end_vertices.shape[0], end_vertices.shape[1]))
if args.tracking is not None:
lines = compress_streamlines(np.swapaxes(flow, 0, 1))
lines_polydata = lines_to_vtk_polydata(lines, None, np.float32)
save_polydata(lines_polydata, args.tracking, True)
# save
if args.end_points is not None:
# save only not masked points
if args.mask is not None:
np.save(args.end_points, end_vertices[mask])
else:
np.save(args.end_points, end_vertices)
if args.end_normal is not None:
end_normal = end_mesh.vertices_normal(args.ed_not_normed, args.ed_not_weighted)
