Brain-age in ultra-low-field MRI: How well does it work?

1. import os
2. import hashlib
3. import re
4. from typing import List, Optional, Union, Callable
6. import pandas as pd
7. import nibabel as nib
8. import torch
9. from torch.utils.data import DataLoader
10. from monai import transforms
11. from monai.data import Dataset
12. from monai.networks.nets.densenet import DenseNet201
13. from nibabel.nifti1 import Nifti1Image
14. from nibabel.nifti2 import Nifti2Image
15. from huggingface_hub import hf_hub_download
17. from .registration import align
18. from .skull_stripping import SkullStripping
21. def _split_nii_name(filename: str) -> tuple:
22. if filename.endswith('.nii.gz'):
23. return filename[:-7], '.nii.gz'
24. return os.path.splitext(filename)
27. def _sanitize_path_component(component: str) -> str:
28. component = re.sub(r'[^A-Za-z0-9._-]+', '_', component).strip('._-')
29. return component or 'path'
32. def _preprocessed_output_paths(input_list: List[str], preprocess_outdir: str) -> List[str]:
33. """
34. Build readable, unique output paths from the shortest distinguishing input suffix.
36. For example, subject_a/visit_a/t1.nii.gz and subject_b/visit_a/t1.nii.gz
37. become subject_a-visit_a-t1.nii.gz and subject_b-visit_a-t1.nii.gz.
38. """
39. if not input_list:
40. return []
42. path_parts = []
43. for input_path in input_list:
44. parts = [p for p in os.path.normpath(input_path).split(os.sep) if p]
45. path_parts.append(parts or [os.path.basename(input_path)])
47. max_depth = max(len(parts) for parts in path_parts)
48. output_names = None
49. for depth in range(1, max_depth + 1):
50. candidate_names = []
51. for parts in path_parts:
52. suffix_parts = parts[-depth:]
53. stem, suffix = _split_nii_name(suffix_parts[-1])
54. name_parts = suffix_parts[:-1] + [stem]
55. candidate_name = '-'.join(_sanitize_path_component(p) for p in name_parts) + suffix
56. candidate_names.append(candidate_name)
58. comparable_names = [name.lower() for name in candidate_names]
59. if len(comparable_names) == len(set(comparable_names)):
60. output_names = candidate_names
61. break
63. if output_names is None:
64. output_names = []
65. for index, (input_path, parts) in enumerate(zip(input_list, path_parts)):
66. stem, suffix = _split_nii_name(parts[-1])
67. digest = hashlib.sha1(f'{index}:{os.path.abspath(input_path)}'.encode()).hexdigest()[:8]
68. output_names.append(f'{_sanitize_path_component(stem)}-{digest}{suffix}')
70. return [os.path.join(preprocess_outdir, name) for name in output_names]
73. class SynthBA:
74. """
75. Predict the brain age from MRI scans of any contrast and resolution using the SynthBA model.
76. """
78. def __init__(
79. self,
80. device: str,
81. checkpoint: Optional[str] = None,
82. model_type: str = 'g',
83. skull_stripping: Optional[SkullStripping] = None,
84. align_fn: Optional[Callable] = None,
85. ):
86. """
87. Initialize the SynthBA brain age prediction model.
89. Args:
90. device (str): The device to run the model on (e.g., 'cuda' or 'cpu').
91. checkpoint (Optional[str]): Path to a model checkpoint. If None, downloads from Hugging Face.
92. model_type (str): Model variant to use ('u' or 'g').
93. skull_stripping (Optional[SkullStripping]): Custom skull stripping method, if provided.
94. align_fn (Optional[Callable]): Custom image alignment function, if provided.
95. """
96. if checkpoint is None and model_type not in ['u', 'g']:
97. raise Exception('SynthBA `model_type` should be either `u` or `g`')
99. self.device = device
100. self.hf_fname = f'synthba-{model_type}.pth'
102. if checkpoint is None:
103. checkpoint = hf_hub_download(repo_id='lemuelpuglisi/synthba', filename=self.hf_fname)
105. checkpoint = torch.load(checkpoint, map_location=device)
106. self.model = DenseNet201(3, 1, 1, dropout_prob=0)
107. self.model.load_state_dict(checkpoint)
108. self.model = self.model.to(device).eval()
110. # Preprocessing components
111. self.skull_stripping = SkullStripping(self.device) if skull_stripping is None else skull_stripping
112. self.align_fn = align if align_fn is None else align_fn
114. # Download alignment templates for different MR weightings
115. self.templates = {
116. 't1': nib.load(hf_hub_download(repo_id='lemuelpuglisi/synthba', filename='MNI152_T1_1mm_Brain.nii.gz')),
117. 't2': nib.load(hf_hub_download(repo_id='lemuelpuglisi/synthba', filename='MNI152_T2_1mm_Brain.nii.gz'))
118. }
120. # Image transformation pipelines
121. self.transforms_fn_direct = transforms.Compose([
122. transforms.EnsureChannelFirst(channel_dim='no_channel'),
123. transforms.Spacing(pixdim=1.4),
124. transforms.ResizeWithPadOrCrop(spatial_size=(130, 130, 130), mode='minimum'),
125. transforms.ScaleIntensity(minv=0, maxv=1),
126. ])
128. self.transforms_fn_multiple = transforms.Compose([
129. transforms.CopyItemsD(keys={'scan_path'}, names=['image']),
130. transforms.LoadImageD(keys='image'),
131. transforms.EnsureChannelFirstD(keys='image', channel_dim='no_channel'),
132. transforms.SpacingD(keys='image', pixdim=1.4),
133. transforms.ResizeWithPadOrCropD(keys='image', spatial_size=(130, 130, 130), mode='minimum'),
134. transforms.ScaleIntensityD(keys='image', minv=0, maxv=1),
135. transforms.Lambda(lambda d: d['image']),
136. ])
139. @torch.inference_mode()
140. def run(
141. self,
142. scan: Union[Nifti1Image, Nifti2Image],
143. preprocess: bool = True,
144. mr_weighting: str = 't1'
145. ) -> float:
146. """
147. Predict brain age for a single MRI scan.
149. Args:
150. scan (Nifti1Image | Nifti2Image): Input MRI scan.
151. preprocess (bool): Whether to apply skull stripping and alignment.
152. mr_weighting (str): MRI contrast type ('t1' or 't2') for template alignment.
154. Returns:
155. float: Predicted brain age in years.
156. """
157. if preprocess:
158. scan = self._preprocess(scan, mr_weighting)
159. scan_tensor = self.transforms_fn_direct(scan.get_fdata())
160. scan_tensor = scan_tensor.unsqueeze(0).to(self.device).float()
161. return self.model(scan_tensor).view(-1).item() * 100.
164. @torch.inference_mode()
165. def run_multiple(
166. self,
167. input_list: List[str],
168. batch_size: int = 1,
169. preprocess: bool = True,
170. preprocess_outdir: Optional[str] = None,
171. mr_weighting: str = 't1',
172. ) -> pd.DataFrame:
173. """
174. Predict brain age for multiple MRI scans in batch.
176. Args:
177. input_list (List[str]): List of paths to input MRI scans.
178. batch_size (int): Number of images processed per batch.
179. preprocess (bool): Whether to apply skull stripping and alignment.
180. preprocess_outdir (Optional[str]): Directory to save preprocessed images.
181. mr_weighting (str): MRI contrast type ('t1' or 't2') for template alignment.
183. Returns:
184. pd.DataFrame: DataFrame with columns ['path', 'pred'] containing predicted brain ages.
185. """
186. if preprocess and preprocess_outdir is None:
187. raise Exception('Please specify where to store the preprocessing output with preprocess_outdir')
189. if preprocess:
190. prep_input_list = []
191. output_paths = _preprocessed_output_paths(input_list, preprocess_outdir)
192. for inp_path, out_path in zip(input_list, output_paths):
193. self._preprocess(nib.load(inp_path), mr_weighting).to_filename(out_path)
194. prep_input_list.append(out_path)
195. input_list = prep_input_list
197. data = [{'scan_path': p} for p in input_list]
198. dataset = Dataset(data=data, transform=self.transforms_fn_multiple)
199. loader = DataLoader(dataset=dataset, batch_size=batch_size)
201. brain_age_list = []
202. for _, images in enumerate(loader):
203. brain_ages = self.model(images.to(self.device))
204. brain_age_list += list(brain_ages.view(-1).cpu().numpy() * 100.)
206. data = [{'path': p, 'pred': y} for p, y in zip(input_list, brain_age_list)]
207. return pd.DataFrame(data)
210. def _preprocess(self, scan: Union[Nifti1Image, Nifti2Image], mr_weighting: str):
211. """
212. Apply skull stripping and spatial alignment to an MRI scan.
214. Args:
215. scan (Nifti1Image | Nifti2Image): Input MRI scan.
216. mr_weighting (str): MRI contrast type ('t1' or 't2') for selecting the reference template.
218. Returns:
219. Nifti1Image | Nifti2Image: Preprocessed and aligned MRI scan.
220. """
221. if mr_weighting not in self.templates.keys():
222. raise Exception('Unknown weighting. Please select between: ' + ', '.join(self.templates.keys()))
223. scan = self.skull_stripping.run(scan)[0]
224. scan = self.align_fn(scan, self.templates[mr_weighting])
225. return scan

synthba.py at commit cc6ca27, no license · at the source