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

1"""Random walker segmentation algorithm. 

2 

3from *Random walks for image segmentation*, Leo Grady, IEEE Trans 

4Pattern Anal Mach Intell. 2006 Nov;28(11):1768-83. 

5 

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""" 

10 

11import warnings 

12 

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 

18 

19from nilearn._utils.helpers import compare_version 

20from nilearn._utils.logger import find_stack_level 

21 

22 

23def _make_graph_edges_3d(n_x, n_y, n_z): 

24 """Return a list of edges for a 3D image. 

25 

26 Parameters 

27 ---------- 

28 n_x : integer 

29 The size of the grid in the x direction. 

30 

31 n_y : integer 

32 The size of the grid in the y direction. 

33 

34 n_z : integer 

35 The size of the grid in the z direction. 

36 

37 Returns 

38 ------- 

39 edges : (2, N) ndarray 

40 With the total number of edges: 

41 

42 N = n_x * n_y * (nz - 1) + 

43 n_x * (n_y - 1) * nz + 

44 (n_x - 1) * n_y * nz 

45 

46 Graph edges with each column describing a node-id pair. 

47 

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 

59 

60 

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 

74 

75 

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] 

81 

82 

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() 

102 

103 

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 

109 

110 

111def _build_ab(lap_sparse, labels): 

112 """Build the matrix A and rhs B of the linear system to solve. 

113 

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 

131 

132 

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 

152 

153 

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 

163 

164 

165def random_walker(data, labels, beta=130, tol=1.0e-3, copy=True, spacing=None): 

166 """Random walker algorithm for segmentation from markers. 

167 

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. 

174 

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. 

182 

183 beta : float, default=130 

184 Penalization coefficient for the random walker motion 

185 (the greater `beta`, the more difficult the diffusion). 

186 

187 tol : float, default=1e-3 

188 Tolerance to achieve when solving the linear system, in 

189 cg' mode. 

190 

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. 

195 

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. 

199 

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. 

206 

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. 

212 

213 The algorithm was first proposed in [1]_. 

214 

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. 

218 

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: 

226 

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 

229 

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. 

232 

233 When the Laplacian is decomposed into blocks of marked and unmarked 

234 pixels:: 

235 

236 L = M B.T 

237 B A 

238 

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:: 

241 

242 A x = - B x_m 

243 

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. 

247 

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. 

252 

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 

263 

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 

270 

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] 

281 

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 ) 

294 

295 if copy: 

296 labels = np.copy(labels) 

297 label_values = np.unique(labels) 

298 

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 

324 

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 

334 

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) 

342 

343 lap_sparse, B = _build_ab(lap_sparse, labels) 

344 

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) 

350 

351 # Clean up results 

352 X = _clean_labels_ar(X + 1, labels).reshape(dims) 

353 return X 

354 

355 

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. 

359 

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 ] 

372 

373 X = np.array(X) 

374 X = np.argmax(X, axis=0) 

375 return X