Coverage for nilearn/_utils/segmentation.py: 12%
126 statements
« prev ^ index » next coverage.py v7.9.1, created at 2025-06-20 10:58 +0200
« prev ^ index » next coverage.py v7.9.1, created at 2025-06-20 10:58 +0200
1"""Random walker segmentation algorithm.
3from *Random walks for image segmentation*, Leo Grady, IEEE Trans
4Pattern Anal Mach Intell. 2006 Nov;28(11):1768-83.
6This code is mostly adapted from scikit-image 0.11.3 release.
7Location of file in scikit image: random_walker function and its supporting
8sub functions in skimage.segmentation
9"""
11import warnings
13import numpy as np
14from scipy import __version__, sparse
15from scipy import ndimage as ndi
16from scipy.sparse.linalg import cg
17from sklearn.utils import as_float_array
19from nilearn._utils.helpers import compare_version
20from nilearn._utils.logger import find_stack_level
23def _make_graph_edges_3d(n_x, n_y, n_z):
24 """Return a list of edges for a 3D image.
26 Parameters
27 ----------
28 n_x : integer
29 The size of the grid in the x direction.
31 n_y : integer
32 The size of the grid in the y direction.
34 n_z : integer
35 The size of the grid in the z direction.
37 Returns
38 -------
39 edges : (2, N) ndarray
40 With the total number of edges:
42 N = n_x * n_y * (nz - 1) +
43 n_x * (n_y - 1) * nz +
44 (n_x - 1) * n_y * nz
46 Graph edges with each column describing a node-id pair.
48 """
49 vertices = np.arange(n_x * n_y * n_z).reshape((n_x, n_y, n_z))
50 edges_deep = np.vstack(
51 (vertices[:, :, :-1].ravel(), vertices[:, :, 1:].ravel())
52 )
53 edges_right = np.vstack(
54 (vertices[:, :-1].ravel(), vertices[:, 1:].ravel())
55 )
56 edges_down = np.vstack((vertices[:-1].ravel(), vertices[1:].ravel()))
57 edges = np.hstack((edges_deep, edges_right, edges_down))
58 return edges
61def _compute_weights_3d(data, spacing, beta=130, eps=1.0e-6):
62 # Weight calculation is main difference in multispectral version
63 # Original gradient**2 replaced with sum of gradients ** 2
64 gradients = sum(
65 _compute_gradients_3d(data[..., channel], spacing) ** 2
66 for channel in range(data.shape[-1])
67 )
68 # All channels considered together in this standard deviation
69 beta /= 10 * data.std()
70 gradients *= beta
71 weights = np.exp(-gradients)
72 weights += eps
73 return weights
76def _compute_gradients_3d(data, spacing):
77 gr_deep = np.abs(data[:, :, :-1] - data[:, :, 1:]).ravel() / spacing[2]
78 gr_right = np.abs(data[:, :-1] - data[:, 1:]).ravel() / spacing[1]
79 gr_down = np.abs(data[:-1] - data[1:]).ravel() / spacing[0]
80 return np.r_[gr_deep, gr_right, gr_down]
83def _make_laplacian_sparse(edges, weights):
84 """Sparse implementation."""
85 pixel_nb = edges.max() + 1
86 diag = np.arange(pixel_nb)
87 i_indices = np.hstack((edges[0], edges[1]))
88 j_indices = np.hstack((edges[1], edges[0]))
89 data = np.hstack((-weights, -weights))
90 lap = sparse.coo_matrix(
91 (data, (i_indices, j_indices)), shape=(pixel_nb, pixel_nb)
92 )
93 connect = -np.ravel(lap.sum(axis=1))
94 lap = sparse.coo_matrix(
95 (
96 np.hstack((data, connect)),
97 (np.hstack((i_indices, diag)), np.hstack((j_indices, diag))),
98 ),
99 shape=(pixel_nb, pixel_nb),
100 )
101 return lap.tocsr()
104def _clean_labels_ar(X, labels):
105 X = X.astype(labels.dtype)
106 labels = np.ravel(labels)
107 labels[labels == 0] = X
108 return labels
111def _build_ab(lap_sparse, labels):
112 """Build the matrix A and rhs B of the linear system to solve.
114 A and B are two block of the laplacian of the image graph.
115 """
116 labels = labels[labels >= 0]
117 indices = np.arange(labels.size)
118 unlabeled_indices = indices[labels == 0]
119 seeds_indices = indices[labels > 0]
120 # The following two lines take most of the time in this function
121 B = lap_sparse[unlabeled_indices][:, seeds_indices]
122 lap_sparse = lap_sparse[unlabeled_indices][:, unlabeled_indices]
123 nlabels = labels.max()
124 rhs = []
125 for lab in range(1, nlabels + 1):
126 mask = labels[seeds_indices] == lab
127 fs = sparse.csr_matrix(mask)
128 fs = fs.transpose()
129 rhs.append(B * fs)
130 return lap_sparse, rhs
133def _mask_edges_weights(edges, weights, mask):
134 """Remove edges of the graph connected to masked nodes, \
135 as well as corresponding weights of the edges.
136 """
137 mask0 = np.hstack(
138 (mask[:, :, :-1].ravel(), mask[:, :-1].ravel(), mask[:-1].ravel())
139 )
140 mask1 = np.hstack(
141 (mask[:, :, 1:].ravel(), mask[:, 1:].ravel(), mask[1:].ravel())
142 )
143 ind_mask = np.logical_and(mask0, mask1)
144 edges, weights = edges[:, ind_mask], weights[ind_mask]
145 max_node_index = edges.max()
146 # Reassign edges labels to 0, 1, ... edges_number - 1
147 order = np.searchsorted(
148 np.unique(edges.ravel()), np.arange(max_node_index + 1)
149 )
150 edges = order[edges.astype(np.int64)]
151 return edges, weights
154def _build_laplacian(data, spacing, mask=None, beta=50):
155 l_x, l_y, l_z = tuple(data.shape[i] for i in range(3))
156 edges = _make_graph_edges_3d(l_x, l_y, l_z)
157 weights = _compute_weights_3d(data, spacing, beta=beta, eps=1.0e-10)
158 if mask is not None:
159 edges, weights = _mask_edges_weights(edges, weights, mask)
160 lap = _make_laplacian_sparse(edges, weights)
161 del edges, weights
162 return lap
165def random_walker(data, labels, beta=130, tol=1.0e-3, copy=True, spacing=None):
166 """Random walker algorithm for segmentation from markers.
168 Parameters
169 ----------
170 data : array_like
171 Image to be segmented in phases.
172 Data spacing is assumed isotropic unless
173 the `spacing` keyword argument is used.
175 labels : array of ints, of same shape as `data` without channels dimension
176 Array of seed markers labeled with different positive integers
177 for different phases. Zero-labeled pixels are unlabeled pixels.
178 Negative labels correspond to inactive pixels that are not taken
179 into account (they are removed from the graph). If labels are not
180 consecutive integers, the labels array will be transformed so that
181 labels are consecutive.
183 beta : float, default=130
184 Penalization coefficient for the random walker motion
185 (the greater `beta`, the more difficult the diffusion).
187 tol : float, default=1e-3
188 Tolerance to achieve when solving the linear system, in
189 cg' mode.
191 copy : bool, default=True
192 If copy is False, the `labels` array will be overwritten with
193 the result of the segmentation. Use copy=False if you want to
194 save on memory.
196 spacing : iterable of floats, optional
197 Spacing between voxels in each spatial dimension. If `None`, then
198 the spacing between pixels/voxels in each dimension is assumed 1.
200 Returns
201 -------
202 output : ndarray
203 An array of ints of same shape as `data`, in which each pixel has
204 been labeled according to the marker that reached the pixel first
205 by anisotropic diffusion.
207 Notes
208 -----
209 The `spacing` argument is specifically for anisotropic datasets, where
210 data points are spaced differently in one or more spatial dimensions.
211 Anisotropic data is commonly encountered in medical imaging.
213 The algorithm was first proposed in [1]_.
215 The algorithm solves the diffusion equation at infinite times for
216 sources placed on markers of each phase in turn. A pixel is labeled with
217 the phase that has the greatest probability to diffuse first to the pixel.
219 The diffusion equation is solved by minimizing x.T L x for each phase,
220 where L is the Laplacian of the weighted graph of the image, and x is
221 the probability that a marker of the given phase arrives first at a pixel
222 by diffusion (x=1 on markers of the phase, x=0 on the other markers, and
223 the other coefficients are looked for). Each pixel is attributed the label
224 for which it has a maximal value of x. The Laplacian L of the image
225 is defined as:
227 - L_ii = d_i, the number of neighbors of pixel i (the degree of i)
228 - L_ij = -w_ij if i and j are adjacent pixels
230 The weight w_ij is a decreasing function of the norm of the local gradient.
231 This ensures that diffusion is easier between pixels of similar values.
233 When the Laplacian is decomposed into blocks of marked and unmarked
234 pixels::
236 L = M B.T
237 B A
239 with first indices corresponding to marked pixels, and then to unmarked
240 pixels, minimizing x.T L x for one phase amount to solving::
242 A x = - B x_m
244 where x_m = 1 on markers of the given phase, and 0 on other markers.
245 This linear system is solved in the algorithm using a direct method for
246 small images, and an iterative method for larger images.
248 References
249 ----------
250 .. [1] Random walks for image segmentation, Leo Grady, IEEE Trans Pattern
251 Anal Mach Intell. 2006 Nov;28(11):1768-83.
253 """
254 out_labels = np.copy(labels)
255 if (labels != 0).all():
256 warnings.warn(
257 "Random walker only segments unlabeled areas, where "
258 "labels == 0. No zero valued areas in labels were "
259 "found. Returning provided labels.",
260 stacklevel=find_stack_level(),
261 )
262 return out_labels
264 if (labels == 0).all():
265 warnings.warn(
266 "Random walker received no seed label. Returning provided labels.",
267 stacklevel=find_stack_level(),
268 )
269 return out_labels
271 # We take multichannel as always False since we are not strictly using
272 # for image processing as such with RGB values.
273 multichannel = False
274 if not multichannel:
275 if data.ndim < 2 or data.ndim > 3:
276 raise ValueError(
277 "For non-multichannel input, data must be of dimension 2 or 3."
278 )
279 dims = data.shape # To reshape final labeled result
280 data = np.atleast_3d(as_float_array(data))[..., np.newaxis]
282 # Spacing kwarg checks
283 if spacing is None:
284 spacing = np.asarray((1.0,) * 3)
285 elif len(spacing) == len(dims):
286 spacing = (
287 np.r_[spacing, 1.0] if len(spacing) == 2 else np.asarray(spacing)
288 )
289 else:
290 raise ValueError(
291 "Input argument `spacing` incorrect, should be an "
292 "iterable with one number per spatial dimension."
293 )
295 if copy:
296 labels = np.copy(labels)
297 label_values = np.unique(labels)
299 # Reorder label values to have consecutive integers (no gaps)
300 if np.any(np.diff(label_values) != 1):
301 mask = labels >= 0
302 labels[mask] = np.searchsorted(
303 np.unique(labels[mask]), labels[mask]
304 ).astype(labels.dtype)
305 labels = labels.astype(np.int32)
306 # If the array has pruned zones, we can have two problematic situations:
307 # - isolated zero-labeled pixels that cannot be determined because they
308 # are not connected to any seed.
309 # - isolated seeds, that is pixels with labels > 0
310 # in connected components without any zero-labeled pixel
311 # to determine.
312 # This causes errors when computing the Laplacian of the graph.
313 # For both cases, the problematic pixels are ignored (label is set to -1).
314 if np.any(labels < 0):
315 # Handle the isolated zero-labeled pixels first
316 filled = ndi.binary_propagation(labels > 0, mask=labels >= 0)
317 labels[np.logical_and(np.logical_not(filled), labels == 0)] = -1
318 del filled
319 # Handle the isolated seeds
320 filled = ndi.binary_propagation(labels == 0, mask=labels >= 0)
321 isolated = np.logical_and(labels > 0, np.logical_not(filled))
322 labels[isolated] = -1
323 del filled
325 # If the operations above yield only -1 pixels
326 if (labels == -1).all():
327 warnings.warn(
328 "Random walker only segments unlabeled areas, where "
329 "labels == 0. Data provided only contains isolated seeds "
330 "and isolated pixels. Returning provided labels.",
331 stacklevel=find_stack_level(),
332 )
333 return out_labels
335 labels = np.atleast_3d(labels)
336 if np.any(labels < 0):
337 lap_sparse = _build_laplacian(
338 data, spacing, mask=labels >= 0, beta=beta
339 )
340 else:
341 lap_sparse = _build_laplacian(data, spacing, beta=beta)
343 lap_sparse, B = _build_ab(lap_sparse, labels)
345 # We solve the linear system
346 # lap_sparse X = B
347 # where X[i, j] is the probability that a marker of label i arrives
348 # first at pixel j by anisotropic diffusion.
349 X = _solve_cg(lap_sparse, B, tol=tol)
351 # Clean up results
352 X = _clean_labels_ar(X + 1, labels).reshape(dims)
353 return X
356def _solve_cg(lap_sparse, B, tol):
357 """Solve lap_sparse X_i = B_i for each phase i, using the conjugate \
358 gradient method.
360 For each pixel, the label i corresponding to the maximal X_i is returned.
361 """
362 lap_sparse = lap_sparse.tocsc()
363 X = [
364 cg(lap_sparse, -b_i.todense(), rtol=tol, atol=0)[0]
365 # TODO
366 # when support scipy to >= 1.12
367 # See https://github.com/nilearn/nilearn/pull/4394
368 if compare_version(__version__, ">=", "1.12")
369 else cg(lap_sparse, -b_i.todense(), tol=tol, atol="legacy")[0]
370 for b_i in B
371 ]
373 X = np.array(X)
374 X = np.argmax(X, axis=0)
375 return X