OSCR

SUITPy: A Python-based toolbox for the analysis of cerebellar functional and anatomical imaging data across the human lifespan.

Code ↔ Paper

6 matches between paragraphs of the paper and lines of its authors' code, computed by the harvester (lexical-v1). Click a colored paragraph or line to see its counterpart.

The 6 matches
  1. [1] § Methods › Cerebellar U-Net isolation › U-Net architecture ↔ SUITPy/isolation.py, lines 554–631 · score 0.78 · corresponding skip connection, transposed convolution, feature maps, concatenated, stride, Block
  2. [2] § Methods › Cerebellar U-Net isolation › Preprocessing ↔ SUITPy/isolation.py, lines 981–1120 · score 0.74 · normalized mutual information, transformation matrix, template space, iteratively, cropped, affine
  3. [3] § Methods › Cerebellar U-Net isolation › U-Net architecture ↔ SUITPy/isolation.py, lines 554–631 · score 0.70 · corresponding skip connection, transposed convolution, convolutional layer, concatenates, strided, Blocks
  4. [4] § Methods › Cerebellar normalization ↔ SUITPy/normalization.py, lines 18–67 · score 0.61 · antsRegistrationSyN, MNI152NLin2009cSymC, space, template, masking, cerebellum
  5. [5] § Methods › Cerebellar U-Net isolation › U-Net architecture ↔ SUITPy/isolation.py, lines 415–493 · score 0.58 · Convolutional Block, convolution layer, padding, kernel, affine, training
  6. [6] § Methods › Cerebellar U-Net isolation › Postprocessing ↔ SUITPy/isolation.py, lines 1242–1259 · score 0.50 · connected component, largest, cluster, voxels, isolation, masks

Paper

Loaded from Europe PMC by your browser, not stored by OSCR: doi.org · Europe PMC

The paper is loaded when this pane is shown.

The authors' code

Python · 1,469 lines · 55 KB · MIT · 5 matches

  1. """
  2. Cerebellar Isolation using a Unet model
  3. @authors: Yao Li, Carlos Hernandez-Castillo, Joern Diedrichsen
  4. """
  6. import sys
  7. import argparse
  8. import os
  9. import nibabel as nib
  10. import ants
  11. import numpy as np
  12. import nitools
  13. from typing import Tuple, Union
  14. import pickle
  15. import warnings
  17. class _Conv3dN:
  18. """
  19. numpy implementation of 3D convolution layer
  20. """
  21. def __init__(
  22. self,
  23. in_channels: int,
  24. out_channels: int,
  25. kernel_size: Union[int, Tuple[int, int, int]],
  26. stride: Union[int, Tuple[int, int, int]] = (1, 1, 1),
  27. padding: Union[int, Tuple[int, int, int]] = (0, 0, 0),
  28. dilation: Union[int, Tuple[int, int, int]] = (1, 1, 1),
  29. bias: bool = True,
  30. padding_mode: str = 'zeros',
  31. ):
  32. self.in_channels = in_channels
  33. self.out_channels = out_channels
  34. self.kernel_size = kernel_size if isinstance(kernel_size, tuple) else (kernel_size,) * 3
  35. self.stride = stride if isinstance(stride, tuple) else (stride,) * 3
  36. self.padding = padding if isinstance(padding, tuple) else (padding,) * 3
  37. self.dilation = dilation if isinstance(dilation, tuple) else (dilation,) * 3
  38. self.bias = bias
  39. self.padding_mode = padding_mode
  41. self.weight = np.random.randn(out_channels, in_channels, *self.kernel_size)
  43. if bias:
  44. self.bias_term = np.zeros(out_channels)
  45. else:
  46. self.bias_term = None
  48. def _pad_input(self, x: np.ndarray) -> np.ndarray:
  49. if all(p == 0 for p in self.padding):
  50. return x
  52. pd, ph, pw = self.padding
  53. if self.padding_mode == 'zeros':
  54. return np.pad(x,
  55. ((0, 0), (0, 0),
  56. (pd, pd), (ph, ph), (pw, pw)),
  57. mode='constant')
  58. elif self.padding_mode == 'reflect':
  59. return np.pad(x,
  60. ((0, 0), (0, 0),
  61. (pd, pd), (ph, ph), (pw, pw)),
  62. mode='reflect')
  63. else:
  64. raise NotImplementedError(f"Padding mode {self.padding_mode} not implemented")
  66. def _im2col(self, x: np.ndarray) -> Tuple[np.ndarray, Tuple[int, int, int]]:
  67. batch_size, in_channels, depth, height, width = x.shape
  68. kd, kh, kw = self.kernel_size
  69. sd, sh, sw = self.stride
  70. pd, ph, pw = self.padding
  71. dd, dh, dw = self.dilation
  73. x_padded = self._pad_input(x)
  75. out_depth = (depth + 2 * pd - dd * (kd - 1) - 1) // sd + 1
  76. out_height = (height + 2 * ph - dh * (kh - 1) - 1) // sh + 1
  77. out_width = (width + 2 * pw - dw * (kw - 1) - 1) // sw + 1
  79. cols = np.zeros((batch_size, in_channels, kd, kh, kw, out_depth, out_height, out_width))
  81. for d in range(kd):
  82. for h in range(kh):
  83. for w in range(kw):
  84. d_start = d * dd
  85. h_start = h * dh
  86. w_start = w * dw
  88. d_slice = slice(d_start, d_start + out_depth * sd, sd)
  89. h_slice = slice(h_start, h_start + out_height * sh, sh)
  90. w_slice = slice(w_start, w_start + out_width * sw, sw)
  92. cols[:, :, d, h, w, :, :, :] = x_padded[:, :, d_slice, h_slice, w_slice]
  94. # Reshape for matrix multiplication (batch, out_d, out_h, out_w, in_ch, kd, kh, kw)
  95. cols = cols.transpose(0, 5, 6, 7, 1, 2, 3, 4)
  96. cols = cols.reshape(batch_size * out_depth * out_height * out_width, -1)
  98. return cols, (out_depth, out_height, out_width)
  100. def load_state(self, weight: np.ndarray, bias_term: np.ndarray):
  101. self.weight = weight
  102. self.bias_term = bias_term
  104. def forward(self, x: np.ndarray) -> np.ndarray:
  105. if x.ndim != 5:
  106. raise ValueError(f"Input must have 5 dimensions (N, C, D, H, W), got {x.ndim}")
  108. batch_size, in_channels, depth, height, width = x.shape
  109. if in_channels != self.in_channels:
  110. raise ValueError(f"Expected {self.in_channels} input channels, got {in_channels}")
  112. cols, (out_depth, out_height, out_width) = self._im2col(x)
  114. # Reshape weights for matrix multiplication (out_ch, in_ch * kd * kh * kw)
  115. weight_flat = self.weight.reshape(self.out_channels, -1)
  117. # Perform convolution via matrix multiplication (batch * out_d * out_h * out_w, out_ch)
  118. output_flat = cols @ weight_flat.T
  120. output = output_flat.reshape(batch_size, out_depth, out_height, out_width, self.out_channels)
  121. output = output.transpose(0, 4, 1, 2, 3) # (batch, out_ch, out_d, out_h, out_w)
  123. if self.bias:
  124. output += self.bias_term.reshape(1, -1, 1, 1, 1)
  126. return output
  128. def __call__(self, x: np.ndarray) -> np.ndarray:
  129. return self.forward(x)
  132. class _ConvTranspose3dN:
  133. """
  134. numpy implementation of 3D transpose convolution layer
  135. """
  136. def __init__(
  137. self,
  138. in_channels: int,
  139. out_channels: int,
  140. kernel_size: Union[int, Tuple[int, int, int]],
  141. stride: Union[int, Tuple[int, int, int]] = (1, 1, 1),
  142. padding: Union[int, Tuple[int, int, int]] = (0, 0, 0),
  143. output_padding: Union[int, Tuple[int, int, int]] = (0, 0, 0),
  144. dilation: Union[int, Tuple[int, int, int]] = (1, 1, 1),
  145. bias: bool = True,
  146. padding_mode: str = "zeros",
  147. ):
  148. self.in_channels = in_channels
  149. self.out_channels = out_channels
  150. self.kernel_size = kernel_size if isinstance(kernel_size, tuple) else (kernel_size,) * 3
  151. self.stride = stride if isinstance(stride, tuple) else (stride,) * 3
  152. self.padding = padding if isinstance(padding, tuple) else (padding,) * 3
  153. self.output_padding = output_padding if isinstance(output_padding, tuple) else (output_padding,) * 3
  154. self.dilation = dilation if isinstance(dilation, tuple) else (dilation,) * 3
  155. self.bias = bias
  156. self.padding_mode = padding_mode
  158. self.weight = np.random.randn(in_channels, out_channels, *self.kernel_size)
  159. if bias:
  160. self.bias_term = np.zeros(out_channels)
  161. else:
  162. self.bias_term = None
  165. def load_state(self, weight: np.ndarray, bias_term: np.ndarray):
  166. self.weight = weight
  167. self.bias_term = bias_term
  169. def _dilate_input(self, x: np.ndarray) -> np.ndarray:
  170. batch_size, in_channels, depth, height, width = x.shape
  171. sd, sh, sw = self.stride
  173. dilated_depth = (depth - 1) * sd + 1
  174. dilated_height = (height - 1) * sh + 1
  175. dilated_width = (width - 1) * sw + 1
  177. dilated_x = np.zeros((batch_size, in_channels, dilated_depth, dilated_height, dilated_width))
  179. # Insert original values at strided positions
  180. dilated_x[:, :, ::sd, ::sh, ::sw] = x
  182. return dilated_x
  184. def _apply_padding(self, x: np.ndarray) -> np.ndarray:
  185. pd, ph, pw = self.padding
  186. kd, kh, kw = self.kernel_size
  188. pad_d_before = kd - 1 - pd
  189. pad_d_after = kd - 1 - pd + self.output_padding[0]
  191. pad_h_before = kh - 1 - ph
  192. pad_h_after = kh - 1 - ph + self.output_padding[1]
  194. pad_w_before = kw - 1 - pw
  195. pad_w_after = kw - 1 - pw + self.output_padding[2]
  197. padding = ((0, 0), (0, 0),
  198. (pad_d_before, pad_d_after),
  199. (pad_h_before, pad_h_after),
  200. (pad_w_before, pad_w_after))
  201. if self.padding_mode == "zeros":
  202. return np.pad(x, padding, mode='constant')
  203. else:
  204. raise NotImplementedError(f"Padding mode {self.padding_mode} not implemented")
  206. def _im2col(self, x: np.ndarray) -> Tuple[np.ndarray, Tuple[int, int, int]]:
  207. batch_size, in_channels, depth, height, width = x.shape
  208. kd, kh, kw = self.kernel_size
  209. sd, sh, sw = self.stride
  210. pd, ph, pw = self.padding
  211. opd, oph, opw = self.output_padding
  212. dd, dh, dw = self.dilation
  214. dilated_x = self._dilate_input(x)
  216. padded_x = self._apply_padding(dilated_x)
  218. # Output dimensions
  219. out_depth = (depth - 1) * sd - 2 * pd + dd * (kd - 1) + opd + 1
  220. out_height = (height - 1) * sh - 2 * ph + dh * (kh - 1) + oph + 1
  221. out_width = (width - 1) * sw - 2 * pw + dw * (kw - 1) + opw + 1
  223. cols = np.zeros((batch_size, in_channels, kd, kh, kw, out_depth, out_height, out_width))
  225. for d in range(kd):
  226. for h in range(kh):
  227. for w in range(kw):
  228. d_start = d * dd
  229. h_start = h * dh
  230. w_start = w * dw
  232. d_slice = slice(d_start, d_start + out_depth)
  233. h_slice = slice(h_start, h_start + out_height)
  234. w_slice = slice(w_start, w_start + out_width)
  236. cols[:, :, d, h, w, :, :, :] = padded_x[:, :, d_slice, h_slice, w_slice]
  238. # Reshape for matrix multiplication (batch, out_d, out_h, out_w, in_ch, kd, kh, kw)
  239. cols = cols.transpose(0, 5, 6, 7, 1, 2, 3, 4)
  240. cols = cols.reshape(batch_size * out_depth * out_height * out_width, -1)
  242. return cols, (out_depth, out_height, out_width)
  244. def forward(self, x: np.ndarray) -> np.ndarray:
  245. if x.ndim != 5:
  246. raise ValueError(f"Input must have 5 dimensions (N, C, D, H, W), got {x.ndim}")
  248. batch_size, in_channels, depth, height, width = x.shape
  249. if in_channels != self.in_channels:
  250. raise ValueError(f"Expected {self.in_channels} input channels, got {in_channels}")
  252. cols, (out_depth, out_height, out_width) = self._im2col(x)
  254. # Reshape weights for matrix multiplication (out_ch, in_ch * kd * kh * kw)
  255. weight = self.weight.transpose(1, 0, 2, 3, 4)
  256. weight_flip = np.flip(weight, (2, 3, 4))
  257. weight_flat = weight_flip.reshape(self.out_channels, -1)
  259. # Perform convolution via matrix multiplication (batch * out_d * out_h * out_w, out_ch)
  260. output_flat = cols @ weight_flat.T
  262. output = output_flat.reshape(batch_size, out_depth, out_height, out_width, self.out_channels)
  263. output = output.transpose(0, 4, 1, 2, 3) # (batch, out_ch, out_d, out_h, out_w)
  265. if self.bias:
  266. output += self.bias_term.reshape(1, -1, 1, 1, 1)
  268. return output
  270. def __call__(self, x: np.ndarray) -> np.ndarray:
  271. return self.forward(x)
  274. class _InstanceNorm3dN:
  275. """
  276. numpy implementation of 3D instance normalization layer
  277. """
  278. def __init__(
  279. self,
  280. num_features: int,
  281. eps: float = 1e-5,
  282. momentum: float = 0.1,
  283. affine: bool = False,
  284. track_running_stats: bool = False
  285. ):
  286. self.num_features = num_features
  287. self.eps = eps
  288. self.momentum = momentum
  289. self.affine = affine
  290. self.track_running_stats = track_running_stats
  292. if self.affine:
  293. self.weight = np.ones(num_features)
  294. self.bias = np.zeros(num_features)
  295. else:
  296. self.weight = None
  297. self.bias = None
  299. if self.track_running_stats:
  300. raise RuntimeError('track_running_stats currently not supported.')
  301. # self.running_mean = np.zeros(num_features)
  302. # self.running_var = np.ones(num_features)
  303. # self.num_batches_tracked = 0
  304. else:
  305. self.running_mean = None
  306. self.running_var = None
  307. self.num_batches_tracked = None
  309. def _check_input_dim(self, x: np.ndarray) -> None:
  310. if x.ndim != 5:
  311. raise ValueError(f"Expected 5D input (got {x.ndim}D input)")
  312. if x.shape[1] != self.num_features:
  313. raise ValueError(f"Expected {self.num_features} channels in input (got {x.shape[1]} channels)")
  315. def _compute_instance_stats(self, x: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
  316. """Compute mean and variance per instance per channel across spatial dimensions (D, H, W)"""
  318. n, c, d, h, w = x.shape
  319. x_reshaped = x.reshape(n, c, -1)
  321. mean = np.mean(x_reshaped, axis=2)
  322. var = np.var(x_reshaped, axis=2)
  324. return mean, var
  326. def _normalize(self, x: np.ndarray, mean: np.ndarray, var: np.ndarray) -> np.ndarray:
  327. """Normalize input using computed statistics"""
  329. n, c, d, h, w = x.shape
  331. mean = mean.reshape(n, c, 1, 1, 1)
  332. var = var.reshape(n, c, 1, 1, 1)
  334. x_normalized = (x - mean) / np.sqrt(var + self.eps)
  336. return x_normalized
  338. def forward(self, x: np.ndarray) -> np.ndarray:
  340. self._check_input_dim(x)
  342. if self.track_running_stats:
  343. n, c, d, h, w = x.shape
  344. mean = self.running_mean.reshape(1, c, 1, 1, 1)
  345. var = self.running_var.reshape(1, c, 1, 1, 1)
  346. else:
  347. # If not tracking running stats, compute instance stats even during inference
  348. mean, var = self._compute_instance_stats(x)
  350. if not self.track_running_stats:
  351. x_normalized = self._normalize(x, mean, var)
  352. else:
  353. x_normalized = (x - mean) / np.sqrt(var + self.eps)
  355. if self.affine:
  356. weight = self.weight.reshape(1, -1, 1, 1, 1)
  357. bias = self.bias.reshape(1, -1, 1, 1, 1)
  358. x_normalized = x_normalized * weight + bias
  360. return x_normalized
  362. def __call__(self, x: np.ndarray) -> np.ndarray:
  363. return self.forward(x)
  366. def load_state_dict(self, state_dict: dict) -> None:
  368. if self.affine:
  369. self.weight = state_dict['weight']
  370. self.bias = state_dict['bias']
  372. if self.track_running_stats:
  373. self.running_mean = state_dict['running_mean']
  374. self.running_var = state_dict['running_var']
  375. self.num_batches_tracked = state_dict['num_batches_tracked']
  378. class _LeakyReLUN:
  379. """
  380. numpy implementation of Leaky ReLUN layer
  381. """
  382. def __init__(self, negative_slope: float = 0.01):
  383. self.negative_slope = negative_slope
  385. def forward(self, x: np.ndarray) -> np.ndarray:
  387. return np.maximum(0, x) + self.negative_slope * np.minimum(0, x)
  389. def __call__(self, x: np.ndarray) -> np.ndarray:
  390. return self.forward(x)
  393. class _SoftmaxN:
  394. """
  395. numpy implementation of Softmax layer
  396. """
  397. def __init__(self, dim: int = -1,):
  398. self.dim = dim
  400. def forward(self, x: np.ndarray) -> np.ndarray:
  401. dim = self.dim
  403. x_shifted = x - np.max(x, axis=dim, keepdims=True)
  405. exp_x = np.exp(x_shifted)
  407. softmax_output = exp_x / np.sum(exp_x, axis=dim, keepdims=True)
  409. return softmax_output
  411. def __call__(self, x: np.ndarray) -> np.ndarray:
  412. return self.forward(x)
  415. class _ConvNetN:
  416. """
  417. convolution block in Unet
  418. """
  419. def __init__(self, in_channels: int, out_channels: int):
  420. """
  422. Args:
  423. in_channels: (int)
  424. the number of input channels
  425. out_channels: (int)
  426. the number of output channels
  427. """
  428. self.conv1 = _Conv3dN(in_channels=in_channels, out_channels=out_channels, kernel_size=3, padding=1)
  429. self.norm1 = _InstanceNorm3dN(num_features=out_channels, affine=True, track_running_stats=False)
  430. self.activate1= _LeakyReLUN()
  432. self.conv2 = _Conv3dN(in_channels=out_channels, out_channels=out_channels, kernel_size=3, padding=1)
  433. self.norm2 = _InstanceNorm3dN(num_features=out_channels, affine=True, track_running_stats=False)
  434. self.activate2 = _LeakyReLUN()
  437. def load_state(
  438. self,
  439. c_weight1: np.ndarray,
  440. c_bias1: np.ndarray,
  441. n_state_dict1: dict,
  442. c_weight2: np.ndarray,
  443. c_bias2: np.ndarray,
  444. n_state_dict2: dict,
  445. ):
  446. """
  447. load the previously trained parameters
  448. Args:
  449. c_weight1: (ndarray)
  450. the weight of the first convolution layer
  451. c_bias1: (ndarray)
  452. the bias of the first convolution layer
  453. n_state_dict1: (dict)
  454. the state dictionary of the first instance normalization layer
  455. c_weight2: (ndarray)
  456. the weight of the second convolution layer
  457. c_bias2: (ndarray)
  458. the bias of the second convolution layer
  459. n_state_dict2: (dict)
  460. the state dictionary of the second instance normalization layer
  462. Returns:
  464. """
  465. self.conv1.load_state(weight=c_weight1, bias_term=c_bias1)
  466. self.norm1.load_state_dict(state_dict=n_state_dict1)
  467. self.conv2.load_state(weight=c_weight2, bias_term=c_bias2)
  468. self.norm2.load_state_dict(state_dict=n_state_dict2)
  470. def forward(self, x: np.ndarray) -> np.ndarray:
  471. """
  472. calculation
  474. Args:
  475. x: (ndarray)
  476. input
  478. Returns:
  479. y: (ndarray)
  480. output
  482. """
  483. result1 = self.conv1(x)
  484. result2 = self.norm1(result1)
  485. result3 = self.activate1(result2)
  486. result4 = self.conv2(result3)
  487. result5 = self.norm2(result4)
  488. result6 = self.activate2(result5)
  490. return result6
  492. def __call__(self, x: np.ndarray) -> np.ndarray:
  493. return self.forward(x)
  496. class _DownSampleN:
  497. """
  498. downsample block in unet
  499. """
  500. def __init__(self, channels: int):
  501. """
  503. Args:
  504. channels: (int)
  505. the number of input channels
  506. """
  507. self.conv1 = _Conv3dN(in_channels=channels, out_channels=channels, kernel_size=3, stride=2, padding=1)
  508. self.norm1 = _InstanceNorm3dN(num_features=channels, affine=True, track_running_stats=False)
  509. self.activate1 = _LeakyReLUN()
  511. def load_state(
  512. self,
  513. c_weight: np.ndarray,
  514. c_bias: np.ndarray,
  515. n_state_dict: dict
  516. ):
  517. """
  518. load the previously trained parameters
  519. Args:
  520. c_weight: (ndarray)
  521. the weight of the convolution layer
  522. c_bias: (ndarray)
  523. the bias of the convolution layer
  524. n_state_dict: (dict)
  525. the state dictionary of the instance normalization layer
  527. """
  528. self.conv1.load_state(weight=c_weight, bias_term=c_bias)
  529. self.norm1.load_state_dict(state_dict=n_state_dict)
  531. def forward(self, x: np.ndarray) -> np.ndarray:
  532. """
  533. calculation
  535. Args:
  536. x: (ndarray)
  537. input
  539. Returns:
  540. y: (ndarray)
  541. output
  543. """
  544. result1 = self.conv1(x)
  545. result2 = self.norm1(result1)
  546. result3 = self.activate1(result2)
  548. return result3
  550. def __call__(self, x: np.ndarray) -> np.ndarray:
  551. return self.forward(x)
  554. class _UpSampleN:
  555. """
  556. upsample block in unet
  557. """
  558. def __init__(self, channels: int):
  559. """
  561. Args:
  562. channels: (int)
  563. the number of input channels
  564. """
  565. self.conv1 = _ConvTranspose3dN(in_channels=channels, out_channels=channels, kernel_size=2, stride=2, padding=0)
  566. self.norm1 = _InstanceNorm3dN(num_features=channels, affine=True, track_running_stats=False)
  567. self.activate1 = _LeakyReLUN()
  569. self.conv2 = _Conv3dN(in_channels=channels * 2, out_channels=channels, kernel_size=3, padding=1)
  570. self.norm2 = _InstanceNorm3dN(num_features=channels, affine=True, track_running_stats=False)
  571. self.activate2 = _LeakyReLUN()
  573. def load_state(
  574. self,
  575. c_weight1: np.ndarray,
  576. c_bias1: np.ndarray,
  577. n_state_dict1: dict,
  578. c_weight2: np.ndarray,
  579. c_bias2: np.ndarray,
  580. n_state_dict2: dict
  581. ):
  582. """
  583. load the previously trained parameters
  584. Args:
  585. c_weight1: (ndarray)
  586. the weight of the transpose convolution layer
  587. c_bias1: (ndarray)
  588. the bias of the transpose convolution layer
  589. n_state_dict1: (dict)
  590. the state dictionary of the first instance normalization layer
  591. c_weight2: (ndarray)
  592. the weight of the convolution layer
  593. c_bias2: (ndarray)
  594. the bias of the convolution layer
  595. n_state_dict2: (dict)
  596. the state dictionary of the second instance normalization layer
  597. Returns:
  599. """
  600. self.conv1.load_state(weight=c_weight1, bias_term=c_bias1)
  601. self.norm1.load_state_dict(state_dict=n_state_dict1)
  602. self.conv2.load_state(weight=c_weight2, bias_term=c_bias2)
  603. self.norm2.load_state_dict(state_dict=n_state_dict2)
  605. def forward(self, x: np.ndarray, feature_map: np.ndarray) -> np.ndarray:
  606. """
  607. calculation
  609. Args:
  610. x: (ndarray)
  611. input
  612. feature_map: (ndarray)
  613. The feature map form the corresponding skip connection
  615. Returns:
  616. y: (ndarray)
  617. output
  619. """
  620. result1 = self.conv1(x)
  621. result2 = self.norm1(result1)
  622. result3 = self.activate1(result2)
  623. result3 = np.concatenate([result3, feature_map], axis=1)
  624. result4 = self.conv2(result3)
  625. result5 = self.norm2(result4)
  626. result6 = self.activate2(result5)
  628. return result6
  630. def __call__(self, x: np.ndarray, feature_map: np.ndarray) -> np.ndarray:
  631. return self.forward(x, feature_map)
  634. class UNetN:
  635. """
  636. Unet model architecture.
  637. """
  639. def __init__(self, init_features: int = 16):
  640. """
  642. Args:
  643. init_features: int
  644. Number of filters in the first convolutional block.
  645. """
  646. self.enc1 = _ConvNetN(in_channels=2, out_channels=init_features)
  647. self.down1 = _DownSampleN(channels=init_features)
  648. self.enc2 = _ConvNetN(in_channels=init_features, out_channels=init_features * 2)
  649. self.down2 = _DownSampleN(init_features * 2)
  650. self.enc3 = _ConvNetN(in_channels=init_features * 2, out_channels=init_features * 4)
  651. self.down3 = _DownSampleN(init_features * 4)
  652. self.enc4 = _ConvNetN(in_channels=init_features * 4, out_channels=init_features * 8)
  653. self.down4 = _DownSampleN(init_features * 8)
  654. self.bottleneck = _ConvNetN(in_channels=init_features * 8, out_channels=init_features * 8)
  655. self.up4 = _UpSampleN(init_features * 8)
  656. self.dec4 = _ConvNetN(in_channels=init_features * 8, out_channels=init_features * 4)
  657. self.up3 = _UpSampleN(init_features * 4)
  658. self.dec3 = _ConvNetN(in_channels=init_features * 4, out_channels=init_features * 2)
  659. self.up2 = _UpSampleN(init_features * 2)
  660. self.dec2 = _ConvNetN(in_channels=init_features * 2, out_channels=init_features)
  661. self.up1 = _UpSampleN(init_features)
  662. self.dec1 = _ConvNetN(in_channels=init_features, out_channels=2)
  663. self.out = _SoftmaxN(dim=1)
  665. def load_state_dict(self, state_dict: dict):
  666. """
  667. Load the state dict of the Unet model.
  668. Args:
  669. state_dict: (dict)
  670. The state dict of the Unet model
  672. """
  674. convs = {'enc1': self.enc1, 'enc2': self.enc2, 'enc3': self.enc3, 'enc4': self.enc4, 'bottleneck': self.bottleneck, 'dec1': self.dec1, 'dec2': self.dec2, 'dec3': self.dec3, 'dec4': self.dec4}
  675. downs = {'down1': self.down1, 'down2': self.down2, 'down3': self.down3, 'down4': self.down4}
  676. ups = {'up1': self.up1, 'up2': self.up2, 'up3': self.up3, 'up4': self.up4}
  678. for c in convs:
  679. c_weight1 = state_dict[f'{c}.layer.0.weight']
  680. c_bias1 = state_dict[f'{c}.layer.0.bias']
  681. n1 = {
  682. 'weight': state_dict[f'{c}.layer.1.weight'],
  683. 'bias': state_dict[f'{c}.layer.1.bias'],
  684. # 'running_mean': state_dict[f'{c}.layer.1.running_mean'],
  685. # 'running_var': state_dict[f'{c}.layer.1.running_var'],
  686. # 'num_batches_tracked': state_dict[f'{c}.layer.1.num_batches_tracked']
  687. }
  688. c_weight2 = state_dict[f'{c}.layer.3.weight']
  689. c_bias2 = state_dict[f'{c}.layer.3.bias']
  690. n2 = {
  691. 'weight': state_dict[f'{c}.layer.4.weight'],
  692. 'bias': state_dict[f'{c}.layer.4.bias'],
  693. # 'running_mean': state_dict[f'{c}.layer.4.running_mean'],
  694. # 'running_var': state_dict[f'{c}.layer.4.running_var'],
  695. # 'num_batches_tracked': state_dict[f'{c}.layer.4.num_batches_tracked']
  696. }
  697. convs[c].load_state(c_weight1=c_weight1, c_bias1=c_bias1, n_state_dict1=n1, c_weight2=c_weight2, c_bias2=c_bias2, n_state_dict2=n2)
  699. for d in downs:
  700. c_weight = state_dict[f'{d}.layer.0.weight']
  701. c_bias = state_dict[f'{d}.layer.0.bias']
  702. n = {
  703. 'weight': state_dict[f'{d}.layer.1.weight'],
  704. 'bias': state_dict[f'{d}.layer.1.bias'],
  705. # 'running_mean': state_dict[f'{d}.layer.1.running_mean'],
  706. # 'running_var': state_dict[f'{d}.layer.1.running_var'],
  707. # 'num_batches_tracked': state_dict[f'{d}.layer.1.num_batches_tracked']
  708. }
  709. downs[d].load_state(c_weight=c_weight, c_bias=c_bias, n_state_dict=n)
  711. for u in ups:
  712. c_weight1 = state_dict[f'{u}.up.0.weight']
  713. c_bias1 = state_dict[f'{u}.up.0.bias']
  714. n1 = {
  715. 'weight': state_dict[f'{u}.up.1.weight'],
  716. 'bias': state_dict[f'{u}.up.1.bias'],
  717. # 'running_mean': state_dict[f'{u}.up.1.running_mean'],
  718. # 'running_var': state_dict[f'{u}.up.1.running_var'],
  719. # 'num_batches_tracked': state_dict[f'{u}.up.1.num_batches_tracked']
  720. }
  721. c_weight2 = state_dict[f'{u}.layer.0.weight']
  722. c_bias2 = state_dict[f'{u}.layer.0.bias']
  723. n2 = {
  724. 'weight': state_dict[f'{u}.layer.1.weight'],
  725. 'bias': state_dict[f'{u}.layer.1.bias'],
  726. # 'running_mean': state_dict[f'{u}.layer.1.running_mean'],
  727. # 'running_var': state_dict[f'{u}.layer.1.running_var'],
  728. # 'num_batches_tracked': state_dict[f'{u}.layer.1.num_batches_tracked']
  729. }
  730. ups[u].load_state(c_weight1=c_weight1, c_bias1=c_bias1, n_state_dict1=n1, c_weight2=c_weight2, c_bias2=c_bias2, n_state_dict2=n2)
  732. def forward(self, t1: np.ndarray, t2: np.ndarray) -> np.ndarray:
  733. """
  734. Unet calculation
  736. Args:
  737. t1: (ndarray)
  738. the cropped T1w image
  739. t2: (ndarray)
  740. the cropped T2w image
  742. Returns:
  743. pm: (ndarray)
  744. the cerebellar probability map
  746. """
  747. result_ini = np.concatenate((t1, t2), axis=1)
  749. result_1 = self.enc1(result_ini)
  750. result_2 = self.enc2(self.down1(result_1))
  751. result_3 = self.enc3(self.down2(result_2))
  752. result_4 = self.enc4(self.down3(result_3))
  753. result_5 = self.bottleneck(self.down4(result_4))
  754. out_1 = self.dec4(self.up4(result_5, result_4))
  755. out_2 = self.dec3(self.up3(out_1, result_3))
  756. out_3 = self.dec2(self.up2(out_2, result_2))
  757. out_4 = self.dec1(self.up1(out_3, result_1))
  759. return self.out(out_4)
  761. def __call__(self, t1: np.ndarray, t2: np.ndarray) -> np.ndarray:
  762. return self.forward(t1, t2)
  764. def _load_model(params_file: str):
  765. """
  766. load model with pretrained weights
  768. Args:
  769. params_file: (string)
  770. path to the pretrained weights
  772. Returns:
  773. net: (Unet)
  774. the pretrained model
  775. """
  776. net = UNetN()
  777. if os.path.exists(params_file):
  778. with open(params_file, "rb") as f:
  779. params = pickle.load(f)
  780. net.load_state_dict(params)
  781. else:
  782. raise RuntimeError('fail to load pre-trained parameters')
  783. return net
  786. def predict(params_file: str, t1: np.ndarray = None, t2: np.ndarray = None) -> np.ndarray:
  787. """
  788. Run a prediction on a single subject using a trained UNet model
  790. Args:
  791. params_file: (string)
  792. filename of the pretrained weights
  793. t1: (ndarray)
  794. Numpy array of T1w cerebellar image (after cropping)
  795. t2: (ndarray)
  796. Numpy array of T2w cerebellar image (after cropping)
  798. Returns:
  799. mask: (ndarray)
  800. the 3D numpy array of predicted mask (template space)
  802. """
  803. net = _load_model(params_file)
  804. if t1 is None:
  805. t1 = np.zeros((128, 128, 128))
  806. else:
  807. t1 = (t1 - t1.mean()) / t1.std()
  808. if t2 is None:
  809. t2 = np.zeros((128, 128, 128))
  810. else:
  811. t2 = (t2 - t2.mean()) / t2.std()
  812. t1, t2 = t1.reshape((1, 1, 128, 128, 128)), t2.reshape((1, 1, 128, 128, 128))
  813. mask = net(t1, t2)
  814. return mask[0][0]
  816. class InputError(Exception):
  817. def __init__(self, message):
  818. super().__init__(message)
  820. def img_read(file: str, use_q_form: bool = False, verbose: bool = False) -> ants.ANTsImage:
  821. """
  822. basic function to read a nifti image as an ANTs image.
  823. checks for consistency of s-form and q-form and
  824. uses s-form by default, or q-form if use_q_form is set to True.
  826. Args:
  827. file: (string)
  828. image path
  829. use_q_form: (bool)
  830. set to True to use q-form
  831. verbose: (bool)
  832. whether to print warning info
  834. Returns:
  835. img (ANTs image)
  836. An ANTs image
  838. """
  840. nib_img = nib.load(file)
  841. s_form, _ = nib_img.get_sform(coded=True)
  842. q_form, _ = nib_img.get_qform(coded=True)
  844. if s_form is None and q_form is None:
  845. raise InputError(f'Both S-form and Q-form are None in {file}')
  846. elif s_form is not None and q_form is None:
  847. if use_q_form:
  848. raise InputError(f'Q-form is None in {file}, Please set use_q_form=False')
  849. else:
  850. nib_img.set_qform(s_form)
  851. elif s_form is None and q_form is not None:
  852. if not use_q_form:
  853. if verbose:
  854. print(f'Warning: S-form is None in {file}. Using Q-form.')
  855. nib_img.set_sform(q_form)
  856. else:
  857. if (s_form != q_form).any():
  858. if not use_q_form:
  859. if verbose:
  860. print(f'Warning: S- and Q-form indicate different image orientations in {file}. Using S-form. Set use_q_form=TRUE to use the q-form.')
  861. nib_img.set_qform(s_form)
  862. else:
  863. if verbose:
  864. print(f'Warning: S- and Q-form indicate different image orientations in {file}. Using Q-form.')
  865. nib_img.set_sform(q_form)
  867. new_img = ants.from_nibabel_nifti(nib_img)
  868. return new_img
  870. def normalized_mutual_information(image1: np.ndarray, image2: np.ndarray, bins:int = 100) -> float:
  871. """
  872. Compute the normalized mutual information between two images
  874. Args:
  875. image1: (ndarray)
  876. First image
  877. image2: (ndarray)
  878. Second image
  879. bins: (int)
  880. number of bins
  882. Returns:
  883. mi: (float)
  884. normalized mutual information between image1 and image2
  886. """
  887. joint_hist, _, _ = np.histogram2d(
  888. image1.ravel(),
  889. image2.ravel(),
  890. bins=bins
  891. )
  893. joint_prob = joint_hist / np.sum(joint_hist)
  895. p_x = np.sum(joint_prob, axis=1)
  896. p_y = np.sum(joint_prob, axis=0)
  898. # H(X) = -Σ p(x) log p(x)
  899. h_x = -1 * np.sum(p_x[p_x > 0] * np.log2(p_x[p_x > 0]))
  901. # H(Y) = -Σ p(y) log p(y)
  902. h_y = -1 * np.sum(p_y[p_y > 0] * np.log2(p_y[p_y > 0]))
  904. # H(X,Y) = -ΣΣ p(x,y) log p(x,y)
  905. h_xy = -1 * np.sum(joint_prob[joint_prob > 0] * np.log2(joint_prob[joint_prob > 0]))
  907. mi = (h_x + h_y) / h_xy
  909. return mi
  912. def registration(img: ants.ANTsImage, brain_mask: ants.ANTsImage = None, template_name: str = 'MNI152NLin6Asym', type_of_transform: str = 'Similarity', max_iterations: int = 5, mi_lower: float = 1.22, mi_upper:float = 1.23) -> Tuple[ants.ANTsTransform, int]:
  913. """
  914. register the image to this template. The function runs ants registration multiple time and find the best registration results evaluated using normalized mutual information.
  916. Args:
  917. img: (ANTsImage)
  918. image to be registered
  919. brain_mask: (ANTsImage)
  920. brain mask of the image
  921. template: (string)
  922. The name of template used. It uses MNI152NLin6Asym by default. (Note that the template must be stored in SUITPy/templates directory and come with brain mask)
  923. type_of_transform: (string)
  924. transform type (Affine by default, check ANTsPY[https://antspy.readthedocs.io/en/latest/registration.html] for details)
  925. max_iterations: (int)
  926. maximum number of iterations
  927. mi_lower: (float)
  928. lower bound of normalized mutual information (range (1, 2)); setting this value too high would result in an error, while setting it too low might cause poor alignment.
  929. mi_upper: (float)
  930. upper bound of normalized mutual information (range (1, 2)); setting this value too high would result in a warning, while setting it too low would prevent the function from searching for better results.
  932. Returns:
  933. trans: (ANTsTransform)
  934. the transformation from the subject space to the template space
  935. status: (int)
  936. the status of the registration (0: for success mi>upper, 1: found one of mi>mi_lower, 2: failure)
  938. """
  939. base_dir = os.path.dirname(__file__)
  940. template = ants.image_read(os.path.join(base_dir, f'templates/tpl-{template_name}_T1w.nii.gz'))
  941. template_brain_mask = ants.image_read(os.path.join(base_dir, f'templates/tpl-{template_name}-brain_mask.nii.gz'))
  942. template_brain = ants.mask_image(template, template_brain_mask)
  943. cnt = 0
  944. status = 2
  945. trans_final = None
  946. best_mi = 0
  948. while cnt < max_iterations:
  950. if brain_mask is not None:
  951. brain = ants.mask_image(img, brain_mask)
  952. result = ants.registration(fixed=template_brain, moving=brain, type_of_transform=type_of_transform)
  953. trans = ants.read_transform(result['fwdtransforms'][0])
  954. else:
  955. result = ants.registration(fixed=template, moving=img, type_of_transform=type_of_transform)
  956. trans = ants.read_transform(result['fwdtransforms'][0])
  958. img_transformed = ants.apply_ants_transform_to_image(trans, img, template)
  959. brain_transformed = ants.mask_image(img_transformed, template_brain_mask)
  960. mi = normalized_mutual_information(brain_transformed.numpy(), template_brain.numpy())
  961. if mi > best_mi:
  962. best_mi = mi
  963. trans_final = trans
  965. if mi > mi_upper:
  966. status = 0
  967. break
  968. if mi > mi_lower:
  969. status = 1
  971. cnt += 1
  973. return trans_final, status
  976. class RegistrationError(Exception):
  977. def __init__(self, message):
  978. super().__init__(message)
  981. class TemplateCerebellarBoundingBox:
  982. """
  983. Basic cerebellar bounding box class, which defines the cropped area.
  984. All other template implementations should be registered to this template.
  985. """
  987. def __init__(self, template_name: str ='MNI152NLin6Asym', bounding_box: np.ndarray = None, cerebellar_center: np.ndarray = None, cropped_size: np.ndarray = None):
  988. """
  989. Create a bounding box
  991. Args:
  992. template_name: (string)
  993. The name of template used. It uses MNI152NLin6Asym by default.
  994. bounding_box: (ndarray)
  995. cerebellar bounding box in MNI space (in mm) (reserved for future development)
  996. cerebellar_center: (ndarray)
  997. (reserved for future development)
  998. cropped_size: (ndarray)
  999. (reserved for future development)
  1000. """
  1001. self.template_name = template_name
  1002. self.cropped_size = (128, 128, 128)
  1003. if bounding_box is not None:
  1004. self.bounding_box = bounding_box
  1005. else:
  1006. if cerebellar_center is not None and cropped_size is not None:
  1007. # reserved for future development
  1008. pass
  1009. else:
  1010. self.bounding_box = np.array([[64, -114, -88], [-64, 14, 40]])
  1012. self.lowerleft = self.bounding_box[0]
  1013. self.upperright = self.bounding_box[1]
  1015. base_dir = os.path.dirname(__file__)
  1016. self.template = ants.image_read(os.path.join(base_dir, f'templates/tpl-{self.template_name}_T1w.nii.gz'))
  1017. self.nib_template = nib.load(os.path.join(base_dir, f'templates/tpl-{self.template_name}_T1w.nii.gz'))
  1018. self.brainmask = ants.image_read(os.path.join(base_dir, f'templates/tpl-{self.template_name}-brain_mask.nii.gz'))
  1019. self.brain = ants.mask_image(self.template, self.brainmask)
  1020. self.affine = nib.load(os.path.join(base_dir, f'templates/tpl-{self.template_name}_T1w.nii.gz')).affine
  1022. def get_crop_indices(self) -> np.ndarray:
  1023. """
  1024. calculate the lower left and upper right indices of the cropped area (in voxels).
  1026. Returns:
  1027. indices: (ndarray)
  1028. a 2 * 3 ndarray consists of two vertices defining the bounding box
  1029. """
  1030. return nitools.coords_to_voxelidxs(self.bounding_box.T, self.nib_template).T
  1032. def get_cropped_affine(self) -> np.ndarray:
  1033. """
  1034. get the cropped area affine
  1036. Returns:
  1037. affine: (ndarray)
  1038. The affine form for the cropped image
  1039. """
  1040. # This function needs to fix. It will fail if the template affine is not diagonal
  1041. affine = np.diag([self.affine[0, 0], self.affine[1, 1], self.affine[2, 2], 1])
  1042. affine[0, 3] = abs(affine[0, 0]) * self.lowerleft[0]
  1043. affine[1, 3] = abs(affine[1, 1]) * self.lowerleft[1]
  1044. affine[2, 3] = abs(affine[2, 2]) * self.lowerleft[2]
  1046. return affine
  1048. def registration(self, img: ants.ANTsImage, brain_mask: ants.ANTsImage = None, type_of_transform: str = 'Similarity', max_iterations: int = 5, mi_lower: float = 1.22, mi_upper:float = 1.23) -> Tuple[ants.ANTsTransform, int]:
  1049. """
  1050. register the image to this template
  1052. Args:
  1053. img: (ANTsImage)
  1054. image to be registered
  1055. brain_mask: (ANTsImage)
  1056. brain mask of the image
  1057. type_of_transform: (string)
  1058. transform type (Affine by default, check ANTsPY[https://antspy.readthedocs.io/en/latest/registration.html] for details)
  1059. max_iterations: (int)
  1060. maximum number of iterations
  1061. mi_lower: (float)
  1062. lower boundary of normalized mutual information
  1063. mi_upper: (float)
  1064. upper boundary of normalized mutual information
  1066. Returns:
  1067. trans: (ANTsTransform)
  1068. the transformation from the subject space to the template space
  1069. status: (int)
  1070. the status of the registration (0 for success, 1 for uncertainty, 2 for failure)
  1072. """
  1074. trans, status = registration(img=img, brain_mask=brain_mask,template_name=self.template_name, type_of_transform=type_of_transform, max_iterations=max_iterations, mi_lower=mi_lower, mi_upper=mi_upper)
  1076. return trans, status
  1078. def crop(self, img: ants.ANTsImage, trans: ants.ANTsTransform = None) -> Tuple[ants.ANTsImage, ants.ANTsImage]:
  1079. """
  1080. Crop the cerebellar area using the bounding box.
  1082. Args:
  1083. img: (ANTsImage)
  1084. image to be cropped
  1085. trans: (ANTsTransform)
  1086. transformation matrix from the image space to the template space (only use it if img is not in the MNI template)
  1088. Returns:
  1089. cropped_img: (ANTsImage)
  1090. cropped image
  1091. img : (ANTsImage)
  1092. the whole transformed image
  1094. """
  1095. start_indices, end_indices = self.get_crop_indices()
  1096. if trans is not None:
  1097. img = ants.apply_ants_transform_to_image(trans, img, self.template)
  1098. return ants.crop_indices(img, tuple(start_indices.astype(int)), tuple(end_indices.astype(int))), img
  1100. def template2subject(self, img: ants.ANTsImage, trans: ants.ANTsTransform, ref: ants.ANTsImage) -> ants.ANTsImage:
  1101. """
  1102. transform the image from template space to the subject space
  1104. Args:
  1105. img: (ANTsImage)
  1106. the image to be transformed
  1107. trans: (ANTs transformation)
  1108. transformation matrix (from subject space to template space)
  1109. ref: (ANTsImage)
  1110. reference image
  1112. Returns:
  1113. img: (ANTsImage)
  1114. the transformed image in subject space
  1116. """
  1117. trans_inv = ants.invert_ants_transform(trans)
  1118. result = ants.apply_ants_transform_to_image(trans_inv, img, ref)
  1120. return result
  1123. def subject_preprocess(t1_file: str = None, t2_file: str = None, brain_mask_file: str = None, label_file: str = None,
  1124. BoundingBox: TemplateCerebellarBoundingBox = TemplateCerebellarBoundingBox(),
  1125. type_of_transform: str = 'Similarity', max_iterations: int = 5, use_q_form: bool = False) -> Tuple[ants.ANTsTransform, ants.ANTsImage, ants.ANTsImage, ants.ANTsImage, ants.ANTsImage, ants.ANTsImage]:
  1126. """
  1127. function to preprocess a single subject.
  1128. 1. Transform the image from subject space to the template space
  1129. 2. Using a pre-defined bounding box to crop the image in the template space
  1131. Args:
  1132. t1_file: (string)
  1133. file to T1w image
  1134. t2_file: (string)
  1135. file to T2w image
  1136. brain_mask_file: (string)
  1137. file to the brain mask image (can be used to improve affine registration)
  1138. label_file: (string)
  1139. file to label image (Optional, this image will be transformed into the template space using the same transformation.)
  1140. BoundingBox: (TemplateCerebellarBoundingBox)
  1141. the bounding box
  1142. type_of_transform: (string)
  1143. reserved for future use (see ANTspy)
  1144. max_iterations: (int)
  1145. maximum number of registration iterations
  1146. use_q_form: (bool)
  1147. set to True to use q-form
  1149. Returns:
  1150. trans: (ANTsTransform)
  1151. transformation from subject space to template space
  1152. t1_crop: (ANTsImage)
  1153. cropped cerebellar area from transformed T1w image
  1154. t2_crop: (ANTsImage)
  1155. cropped cerebellar area from transformed T2w image
  1156. label_crop: (ANTsImage)
  1157. cropped cerebellar area from transformed label image
  1158. t1_whole: (ANTsImage)
  1159. whole transformed T1w image
  1160. t2_whole: (ANTsImage)
  1161. whole transformed T2w image
  1163. """
  1165. if t1_file is not None:
  1166. t1 = img_read(file=t1_file, use_q_form=use_q_form, verbose=True)
  1167. else:
  1168. t1 = None
  1169. if t2_file is not None:
  1170. t2 = img_read(file=t2_file, use_q_form=use_q_form, verbose=True)
  1171. else:
  1172. t2 = None
  1174. if brain_mask_file is not None:
  1175. brain_mask = img_read(file=brain_mask_file, use_q_form=use_q_form)
  1176. else:
  1177. brain_mask = None
  1179. # Read additional images
  1180. if label_file is not None:
  1181. label = img_read(file=label_file, use_q_form=use_q_form)
  1182. else:
  1183. label = None
  1185. # If T1 and T2 are both given, but not in the same spacing, reslice T2 to T1 first
  1186. if t2 is not None and t1 is not None:
  1187. if ants.get_spacing(t1) != ants.get_spacing(t2):
  1188. t2 = ants.registration(fixed=t1, moving=t2, type_of_transform='Rigid')['warpedmovout']
  1190. if t1 is not None:
  1191. trans, status = BoundingBox.registration(img=t1, brain_mask=brain_mask, type_of_transform=type_of_transform, max_iterations=max_iterations)
  1192. else:
  1193. trans, status = BoundingBox.registration(img=t2, brain_mask=brain_mask, type_of_transform=type_of_transform, max_iterations=max_iterations)
  1195. if status == 2:
  1196. raise RegistrationError('invalid registration detected')
  1198. if status == 1:
  1199. warnings.warn(f'low-quality registration detected, please double check the results for {t1_file if t1 is not None else t2_file}')
  1201. if t1 is not None:
  1202. t1_crop, t1_whole = BoundingBox.crop(t1, trans)
  1203. else:
  1204. t1_crop = None
  1205. t1_whole = None
  1207. if t2 is not None:
  1208. t2_crop, t2_whole = BoundingBox.crop(t2, trans)
  1209. else:
  1210. t2_crop = None
  1211. t2_whole = None
  1213. if label is not None:
  1214. label_crop, _ = BoundingBox.crop(label, trans)
  1215. else:
  1216. label_crop = None
  1218. return trans, t1_crop, t2_crop, label_crop, t1_whole, t2_whole
  1221. def threshold(img: ants.ANTsImage, lower: float = 0.5, upper: float = 1.0) -> ants.ANTsImage:
  1222. """
  1223. remove all other values from the image
  1225. Args:
  1226. img: (ANTsImage)
  1227. the input image
  1228. lower: (float)
  1229. lower threshold
  1230. upper: (float)
  1231. upper threshold
  1233. Returns:
  1234. image : (ANTsImage)
  1235. the thresholded image
  1236. """
  1237. img[img < lower] = 0
  1238. img[img > upper] = 0
  1239. return img
  1242. def remove_islands(img: ants.ANTsImage) -> ants.ANTsImage:
  1243. """ Removes parts of the mask that is not connected to the largest cluster
  1245. Args:
  1246. img (ANTsImage): the input image
  1247. Returns:
  1248. mask (ANTsImage): Image containing the largest connected component
  1249. """
  1250. clusters = ants.image_to_cluster_images(img)
  1252. mask = None
  1253. voxels = 0
  1254. for temp in clusters:
  1255. if temp.numpy().sum() > voxels:
  1256. mask = temp
  1257. voxels = temp.numpy().sum()
  1259. return mask
  1262. def subject_postprocess(mask: ants.ANTsImage, trans: ants.ANTsTransform, BoundingBox: TemplateCerebellarBoundingBox, ref: ants.ANTsImage) -> ants.ANTsImage:
  1263. """
  1264. transform the predicted cerebellum mask to the original space
  1265. Args:
  1266. mask: (ANTsImage)
  1267. the predicted cerebellum mask from the template space
  1268. trans: (ANTsTransform)
  1269. the transformation from subject space to template space
  1270. BoundingBox: (TemplateCerebellarBoundingBox)
  1271. the bounding box
  1272. ref: (ANTsImage)
  1273. the reference image
  1275. Returns:
  1276. result: (ANTsImage)
  1277. the final cerebellum mask from the subject space
  1279. """
  1281. result = BoundingBox.template2subject(mask, trans, ref)
  1282. # threshold and binarize the image
  1283. result = threshold(result)
  1284. result[result != 0] = 1
  1286. result = remove_islands(result)
  1287. return result
  1290. def isolate(t1_file: str = None, t2_file: str = None,
  1291. brain_mask_file: str = None,
  1292. label_file: str = None,
  1293. result_folder: str = None,
  1294. template: str = 'MNI152NLin6Asym',
  1295. type_of_transform: str = 'Similarity',
  1296. max_iterations: int = 5,
  1297. params: str = 'pre_trained_numpy.pkl',
  1298. save_cropped_files: bool = False,
  1299. use_q_form: bool = False,
  1300. verbose: bool = True) -> ants.ANTsImage:
  1301. """
  1302. main function for cerebellum isolation
  1304. Args:
  1305. t1_file: (string)
  1306. filename and path to T1w image, optional
  1307. t2_file: (string)
  1308. filename and path to T2w image, optional
  1309. brain_mask_file: (string)
  1310. filename and path to brain mask, optional
  1311. label_file: (string)
  1312. filename and path to label image, optional (reserved, currently has no effect)
  1313. result_folder: (string)
  1314. path to output folder (optional, otherwise it is saved to input image folder)
  1315. template: (string)
  1316. template to use (reserved)
  1317. type_of_transform: (string)
  1318. reserved for future use (see ANTspy)
  1319. max_iterations: (int)
  1320. maximum number of registration iterations (optional, default 5)
  1321. params: (string)
  1322. path to params file (reserved)
  1323. save_cropped_files: (bool)
  1324. set to True to save files cropped to window
  1325. use_q_form: (bool)
  1326. set to True to use q-form
  1327. verbose: (bool)
  1328. whether to print out status information during processing
  1330. """
  1332. if t1_file is not None:
  1333. result_folder = os.path.dirname(os.path.abspath(t1_file)) if result_folder is None else result_folder
  1334. basename = os.path.splitext(os.path.basename(t1_file))
  1335. elif t2_file is not None:
  1336. result_folder = os.path.dirname(os.path.abspath(t2_file)) if result_folder is None else result_folder
  1337. basename = os.path.splitext(os.path.basename(t2_file))
  1338. else:
  1339. raise RuntimeError('Must specify either t1_file or t2_file')
  1341. # Strip .nii or .nii.gz extension
  1342. if basename[1] == '.gz':
  1343. basename = os.path.splitext(basename[0])
  1344. basename = basename[0]
  1346. # find paramter file and template bounding box
  1347. base_dir = os.path.dirname(os.path.abspath(__file__))
  1348. params_file = os.path.join(base_dir, 'parameters', params)
  1349. BoundingBox = TemplateCerebellarBoundingBox(template_name=template)
  1351. try:
  1352. # Crop the images to the Unet input window
  1353. if verbose:
  1354. print(f"preprocessing {t1_file if t1_file is not None else t2_file}")
  1355. trans, t1_crop, t2_crop, label_crop, _, _ = subject_preprocess(t1_file=t1_file,
  1356. t2_file=t2_file,
  1357. brain_mask_file=brain_mask_file,
  1358. label_file=label_file,
  1359. BoundingBox=BoundingBox,
  1360. type_of_transform=type_of_transform,
  1361. max_iterations=max_iterations,
  1362. use_q_form=use_q_form)
  1363. if isinstance(t1_crop, ants.core.ants_image.ANTsImage):
  1364. t1_crop_data = t1_crop.numpy()
  1365. else:
  1366. t1_crop_data = t1_crop
  1367. if isinstance(t2_crop, ants.core.ants_image.ANTsImage):
  1368. t2_crop_data = t2_crop.numpy()
  1369. else:
  1370. t2_crop_data = t2_crop
  1371. if isinstance(label_crop, ants.core.ants_image.ANTsImage):
  1372. label_crop_data = label_crop.numpy()
  1373. else:
  1374. label_crop_data = label_crop
  1377. # Do a forward pass through the Unet model
  1378. if verbose:
  1379. print('isolating cerebellum using UNet model')
  1380. mask_template = predict(params_file=params_file, t1=t1_crop_data, t2=t2_crop_data)
  1381. mask_template = nib.Nifti1Image(mask_template, BoundingBox.get_cropped_affine())
  1382. mask_template = ants.from_nibabel_nifti(mask_template)
  1384. # Postprocess and transform the mask back to subject space
  1385. if verbose:
  1386. print('postprocessing')
  1387. if t1_file is not None:
  1388. mask_subject = subject_postprocess(mask=mask_template, trans=trans, BoundingBox=BoundingBox, ref=img_read(file=t1_file, use_q_form=use_q_form))
  1389. else:
  1390. mask_subject = subject_postprocess(mask=mask_template, trans=trans, BoundingBox=BoundingBox, ref=img_read(file=t2_file, use_q_form=use_q_form))
  1392. # use the original header info for the dseg mask
  1393. if t1_file is not None:
  1394. ref = nib.load(t1_file)
  1395. else:
  1396. ref = nib.load(t2_file)
  1397. mask_subject = nib.Nifti1Image(mask_subject.numpy().astype(int), affine=ref.affine, header=ref.header)
  1399. os.makedirs(result_folder, exist_ok=True)
  1400. ofname = f'{basename}_cerebellum_dseg.nii.gz'
  1401. if verbose:
  1402. print(f"saving results to {ofname}")
  1403. nib.save(mask_subject, os.path.join(result_folder, ofname))
  1405. if save_cropped_files:
  1406. if verbose:
  1407. print(f"saving intermediate results to {result_folder}")
  1408. if t1_crop is not None:
  1409. ants.image_write(t1_crop, os.path.join(result_folder, f'{basename}_crop.nii.gz'))
  1410. else:
  1411. ants.image_write(t2_crop, os.path.join(result_folder, f'{basename}_crop.nii.gz'))
  1412. ants.image_write(mask_template, os.path.join(result_folder, f'{basename}_cerebellum_crop_dseg.nii.gz'))
  1413. ants.write_transform(trans, os.path.join(result_folder, f'{basename}_trans.mat'))
  1414. except RegistrationError as e:
  1415. print(f'Caught registration error for {t1_file if t1_file is not None else t2_file} : {e}')
  1416. print(f'Isolation fails on {t1_file if t1_file is not None else t2_file}. No results were saved')
  1417. mask_subject = None
  1418. except InputError as e:
  1419. print(f'Caught input file error for {t1_file if t1_file is not None else t2_file} : {e}')
  1420. print(f'Isolation fails on {t1_file if t1_file is not None else t2_file}. No results were saved')
  1421. mask_subject = None
  1422. return img_read(file=os.path.join(result_folder, ofname), use_q_form=use_q_form) if mask_subject is not None else None
  1425. if __name__ == '__main__':
  1426. parser = argparse.ArgumentParser()
  1428. parser.add_argument('--T1', type=str, help='path to T1w image')
  1429. parser.add_argument('--T2', type=str, help='path to T2w image')
  1430. parser.add_argument('--brain_mask', type=str, help='path to brain mask image')
  1431. parser.add_argument('--label', type=str, help='path to label image (reserved, currently has no effect)')
  1432. parser.add_argument('--result_folder', type=str, help='path to save the isolation image (results will be saved to '
  1433. 'T1w image folder (or T2w image folder if no T1w image is '
  1434. 'specified))')
  1435. parser.add_argument('--template', type=str, default='MNI152NLin6Asym',
  1436. help='template for registration (MNI152NLin6Asym by '
  1437. 'default)')
  1438. parser.add_argument('--type_of_transform', type=str, default='Similarity', help='reserved for future use (see ANTspy)')
  1439. parser.add_argument('--max_iterations', type=int, default=5, help='maximum number of registration iterations (optional, default 5)')
  1440. parser.add_argument('--params', type=str, default='pre_trained.pkl', help='pretrained parameter file')
  1441. parser.add_argument('--save_cropped_files', action='store_true', help='whether to save files cropped to UNet input window')
  1442. parser.add_argument('--use_q_form', action='store_true', help='whether to use q-form')
  1443. parser.add_argument('--verbose', action='store_true', help='whether to print out status information during processing')
  1445. args = parser.parse_args()
  1446. print(args)
  1448. if args.T1 is None and args.T2 is None:
  1449. raise RuntimeError('Must specify either t1_file or t2_file')
  1451. if args.result_folder is None:
  1452. if args.T1 is None:
  1453. args.result_folder = os.path.dirname(os.path.abspath(args.T2))
  1454. else:
  1455. args.result_folder = os.path.dirname(os.path.abspath(args.T1))
  1458. isolate(t1_file=args.T1,
  1459. t2_file=args.T2,
  1460. brain_mask_file=args.brain_mask,
  1461. label_file=args.label,
  1462. result_folder=args.result_folder,
  1463. template=args.template,
  1464. type_of_transform=args.type_of_transform,
  1465. max_iterations=args.max_iterations,
  1466. params=args.params,
  1467. save_cropped_files=args.save_cropped_files,
  1468. use_q_form=args.use_q_form,
  1469. verbose=args.verbose)

isolation.py at commit 3d20055, under MIT · at the source

Overview

Authors: Yaping Wang1,2,3, Yao Li4, Bassel Arafat1,2, Vahid Ashkanichenarlogh1,2, Caroline Nettekoven1,2,5, Ana Luísa Pinho1,2,6, Carlos R. Hernandez-Castillo4, Andre F. Marquand3, Jörn Diedrichsen1,2
ORCID iDs: Bassel Arafat, Caroline Nettekoven, Ana Luísa Pinho, Jörn Diedrichsen
  1. Western Institute for Neuroscience, Western University, London, ON, Canada
  2. Department of Computer Science, Western University, London, ON, Canada
  3. Donders Institute for Brain, Cognition and Behaviour, Radboud University Medical Centre, Nijmegen, The Netherlands
  4. Faculty of Computer Science, Dalhousie University, Halifax, Canada
  5. Department of Experimental Psychology, University of Oxford, Oxford, United Kingdom
  6. Department of Psychology, Western University, London, ON, Canada
Institutions: Western University (Canada); Radboud University Nijmegen (Netherlands); Radboud University Medical Center (Netherlands); Donders Institute for Brain, Cognition and Behaviour (Netherlands); Dalhousie University (Canada); University of Oxford (United Kingdom)
Journal: Imaging neuroscience (Cambridge, Mass.), volume 4, article IMAG.a.1323
Dates: received 21 May 2026; accepted 24 June 2026; published online 6 August 2026
Type: Research article · Language: English
License: CC BY
Identifiers: DOI 10.1162/imag.a.1323 · PMID 42569404 · PMCID PMC13449926 · OpenAlex W7167902735
Open access: diamond, a free copy (OpenAlex)
Status: code verified
Categories: human (organism), methods / tools (subfield)
Methods: Connectivity, Statistics, Machine learning, fMRI & imaging
Keywords: human cerebellum, MRI, deep learning, cerebellar isolation, spatial normalization, voxel-wise inference, neuroimaging toolbox, lifespan
Journal subjects: Software Toolbox
Topic: Vestibular and auditory disorders (Neurology, Neuroscience), according to OpenAlex
Funding: Raynor Cerebellum Project; Canadian Institutes of Health Research (PJT-191815); Canada First Research Excellence Fund (BrainsCAN); Wellcome Trust Early Career Award (306553/Z/23/Z); Junior Research Fellowship Grant from Linacre College, University of Oxford
Citations: not cited yet (Europe PMC); 57 references in the paper

Abstract

The human cerebellum plays a central role in motor, emotional, and cognitive functions, and is implicated in many brain disorders. To improve the analysis of functional and anatomical imaging from the cerebellum, we introduce SUITPy, an improved and fully revised Python implementation of the widely used SUIT toolbox. For this new version, we developed a U-Net based model to automatically isolate the cerebellum from adjacent cortical tissue, which achieves higher fidelity than existing algorithms. The isolation works robustly without manual corrections for imaging data across the lifespan. We show that isolation and subsequent normalization to a cerebellum-only template lead to a more precise alignment of cerebellar structures across participants compared to normalization using a whole-brain template. We also show the utility of the cerebellar mask to prevent contamination of cerebellar functional data from surrounding cortical structures. The toolbox also provides functionality for visualizing cerebellar data on a flatmap, along with a range of anatomical and functional cerebellar atlases, thereby offering an essential tool that enables accurate cerebellar analysis across the lifespan.

Reproduced under the paper's license (CC BY), from the paper cited above.

Repository

Its files are read in the Code ↔ Paper reader above, with 6 matches between paragraphs and lines of code.

DiedrichsenLab/SUITPy

License: MIT
State: the link answers, verified on 27 September 2026
Evidence: files inventoried
Commit: 3d20055022e5339687f412da92979520c7568da2, 12 September 2026
Languages: Python (16), Jupyter (8)
Size: 95 files, 24 scripts
Software Heritage: not archived
Found in: “Data and Code Availability”
Holds: README, license file, CITATION.cff, environment (environment.yml, requirements-build-docs.txt, requirements-dev.txt, requirements-min.txt, setup.cfg, setup.py), tests, continuous integration, documentation, 8 notebooks
Tools: NiBabel (17 files), NumPy (11 files), Matplotlib (8 files), Nilearn (8 files), ANTs (4 files), pandas (3 files), Plotly (2 files), SciPy (2 files)
Availability: 1 check, the latest on 27 September 2026: the link answers
  • 27 September 2026: the link answers
26 files
At the source: github.com/DiedrichsenLab/SUITPy

The paper's code and data availability statement is in the Data section.

Tracing map

Proposed by the machine: these links were found in the paper and verified at the source, without human review. The map will receive a Zenodo DOI once one of the paper's authors has validated it with their ORCID.

What the map holds:

  • 1 repository of the authors' code, each at its verified commit, with its license and how the link was found in the paper;
  • 24 scripts, each with its path and the digest of its content;
  • 6 matches between paragraphs of the paper and lines of the code (method lexical-v1);
  • neither the text of the paper nor the code itself.

Its JSON (tracing-map.json) is deposited on Zenodo with its DOI once the map is validated.

Data

No dataset and no data link were found in the paper.

Data and Code Availability

The manual isolation dataset (except for the EssCer dataset) is available upon request. The SUITPy toolbox (version 2.1) can be installed via pip with source code available at https://github.com/DiedrichsenLab/SUITPy. Tutorials and documentation can be found at https://suitpy.readthedocs.io.

Reproduced under the paper's license (CC BY), from the paper cited above.

Versions

The history of this record: each version stored by the harvester or made by a correction of its authors or of the maintainers of its code, and what changed in its facts. The texts of the paper (its abstract, its availability statements) are not part of it; versions that changed only those are not listed.

Version 1, 27 September 2026: the first record

Recorded: type, language, journal, volume, pages, dates, 9 authors, 8 keywords, 5 funders, 55 references.

Cite

This paper

Wang, Y., Li, Y., Arafat, B., Ashkanichenarlogh, V., Nettekoven, C., Pinho, A. L., Hernandez-Castillo, C. R., Marquand, A. F., & Diedrichsen, J. (2026). SUITPy: A Python-based toolbox for the analysis of cerebellar functional and anatomical imaging data across the human lifespan. Imaging neuroscience (Cambridge, Mass.), 4, IMAG.a.1323. https://doi.org/10.1162/imag.a.1323

BibTeX

@article{wang2026suitpy,
author = {Wang, Yaping and Li, Yao and Arafat, Bassel and Ashkanichenarlogh, Vahid and Nettekoven, Caroline and Pinho, Ana Luísa and Hernandez-Castillo, Carlos R. and Marquand, Andre F. and Diedrichsen, Jörn},
title = {{SUITPy: A Python-based toolbox for the analysis of cerebellar functional and anatomical imaging data across the human lifespan}},
journal = {Imaging neuroscience (Cambridge, Mass.)},
year = {2026},
month = aug,
volume = {4},
pages = {IMAG.a.1323},
publisher = {MIT Press},
issn = {2837-6056},
doi = {10.1162/imag.a.1323},
url = {https://doi.org/10.1162/imag.a.1323},
pmid = {42569404},
pmcid = {PMC13449926}
}

RIS

TY - JOUR
AU - Wang, Yaping
AU - Li, Yao
AU - Arafat, Bassel
AU - Ashkanichenarlogh, Vahid
AU - Nettekoven, Caroline
AU - Pinho, Ana Luísa
AU - Hernandez-Castillo, Carlos R.
AU - Marquand, Andre F.
AU - Diedrichsen, Jörn
TI - SUITPy: A Python-based toolbox for the analysis of cerebellar functional and anatomical imaging data across the human lifespan
T2 - Imaging neuroscience (Cambridge, Mass.)
J2 - Imaging Neurosci (Camb)
PY - 2026
DA - 2026/08/06
VL - 4
SP - IMAG.a.1323
SN - 2837-6056
PB - MIT Press
DO - 10.1162/imag.a.1323
UR - https://doi.org/10.1162/imag.a.1323
LA - en
ER -

CSL-JSON

{
"id": "10.1162/imag.a.1323",
"type": "article-journal",
"title": "SUITPy: A Python-based toolbox for the analysis of cerebellar functional and anatomical imaging data across the human lifespan",
"container-title": "Imaging neuroscience (Cambridge, Mass.)",
"author": [
{
"family": "Wang",
"given": "Yaping"
},
{
"family": "Li",
"given": "Yao"
},
{
"family": "Arafat",
"given": "Bassel"
},
{
"family": "Ashkanichenarlogh",
"given": "Vahid"
},
{
"family": "Nettekoven",
"given": "Caroline"
},
{
"family": "Pinho",
"given": "Ana Luísa"
},
{
"family": "Hernandez-Castillo",
"given": "Carlos R."
},
{
"family": "Marquand",
"given": "Andre F."
},
{
"family": "Diedrichsen",
"given": "Jörn"
}
],
"container-title-short": "Imaging Neurosci (Camb)",
"volume": "4",
"page": "IMAG.a.1323",
"DOI": "10.1162/imag.a.1323",
"PMID": "42569404",
"PMCID": "PMC13449926",
"ISSN": "2837-6056",
"publisher": "MIT Press",
"URL": "https://doi.org/10.1162/imag.a.1323",
"language": "en",
"issued": {
"date-parts": [
[
2026,
8,
6
]
]
}
}

The tracing map gets a citation of its own once an author has validated it and it has a DOI.

Similar papers

The papers with a page that share the most with this one: the tools found in their code, their categories, datasets, cited references and authors, the rarest counting most.

[1] doi:10.21203/rs.3.rs-9326213/v1 [code]
Multi-task fMRI outperforms resting-state fMRI for revealing task-invariant organization of the human brain
Journal: Research Square (preprint)
In common: ANTs, Nilearn, Plotly, 5 other tools, 3 references, 3 authors
[2] doi:10.1038/s41467-026-72940-5 [code]
Cerebellar growth is associated with domain-specific cerebral maturation and socio-linguistic behavior.
Journal: Nature communications
In common: NiBabel, pandas, SciPy, 2 other tools, 9 references, author Jörn Diedrichsen
[3] doi:10.1016/j.neuroimage.2026.121930
Quantifying cerebellar signal detectability in MEG and EEG in epilepsy using anatomically informed source modeling.
Journal: NeuroImage
In common: 9 references
[4] doi:10.1186/s12888-026-08071-4 [code]
Cerebellar gray matter volume difference in first-episode bipolar and unipolar depression.
Journal: BMC psychiatry
In common: 8 references
[5] doi:10.1038/s41467-026-72931-6 [code]
Three parsimonious spatiotemporal patterns in cerebellum reveal individual traits in function and behavior.
Journal: Nature communications
In common: Nilearn, NiBabel, pandas, 3 other tools, 4 references
[6] doi:10.1162/imag.a.1276 [code]
High-resolution whole-brain magnetic resonance spectroscopic imaging in youth at risk for psychosis.
Journal: Imaging neuroscience (Cambridge, Mass.)
In common: ANTs, Nilearn, Plotly, 5 other tools, 1 reference
[7] doi:10.1162/imag.a.1362 [code]
Human fMRI at 11.7T: Assessing feasibility, stability, and reliability on the Iseult scanner.
Journal: Imaging neuroscience (Cambridge, Mass.)
In common: ANTs, Nilearn, NiBabel, 4 other tools, methods / tools, 2 references
[8] doi:10.1038/s41592-026-03159-x [code]
Siibra: a software tool suite for realizing a Multilevel Human Brain Atlas from complex data resources.
Journal: Nature methods
In common: Nilearn, Plotly, NiBabel, 4 other tools, methods / tools, 2 references
[9] doi:10.1038/s41467-026-71151-2 [code]
Common and distinct neural correlates of social interaction processing and theory of mind in narratives.
Journal: Nature communications
In common: ANTs, Nilearn, NiBabel, 4 other tools, 2 references
[10] doi:10.1162/imag.a.1269 [code]
From early to contemporary normative modeling: Mapping individual differences in neurophysiological signals.
Journal: Imaging neuroscience (Cambridge, Mass.)
In common: Nilearn, Plotly, NiBabel, 4 other tools, 2 references

Contribute

The authors of this paper can claim it, correct its record and validate its tracing map, and the maintainers of its code (its owner, or a public member of its organization) correct what it says of their repository; anyone signed in can ask for its removal. Every request goes to OSCR's own machine, which answers it; your account page follows them.

Sign in with ORCID to claim this paper as one of its authors, correct its record or validate its tracing map: when the paper's metadata lists your ORCID iD, you are recognized at once. Maintainers of its code: sign in with GitHub, then claim the repository on your account page.

Request its removal

To ask OSCR to remove this record, the copies of its authors' scripts or its tracing map, use the removal request page: signed in, you say who you are, what to remove and why, then review and confirm the request. Published rules decide every request (how).

Discussion, reproductions, activity

Discussion: questions and error reports about this paper and its code, from signed-in readers and its authors. It opens with sign-in.

Reproductions: reports from readers who ran the authors' code: what they reproduced, with which environment, commit and data. It opens with sign-in.

Activity: what happens around this paper: new versions of its record, its map's validation, discussions and reproductions. It opens with sign-in.