Skip to content

Commit ee200fa

Browse files
fepegarfepegar
committed
Add datasets
1 parent 4ececea commit ee200fa

7 files changed

Lines changed: 331 additions & 21 deletions

File tree

‎docs/index.md‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -83,3 +83,14 @@ If you use this library for your research, please cite our paper:
8383
> augmentation and patch-based sampling of medical images in deep learning.*
8484
> Computer Methods and Programs in Biomedicine (June 2021), p. 106236.
8585
> [doi:10.1016/j.cmpb.2021.106236](https://doi.org/10.1016/j.cmpb.2021.106236).
86+
87+
## Related projects
88+
89+
- [MONAI](https://monai.readthedocs.io)
90+
- [Cornucopia](https://cornucopia.readthedocs.io/)
91+
- [batchgenerators](https://github.com/MIC-DKFZ/batchgenerators)[[v2](https://github.com/MIC-DKFZ/batchgeneratorsv2)]
92+
- [volumentations](https://github.com/ZFTurbo/volumentations) (low activity)
93+
- [Rising](https://rising.readthedocs.io) (archived)
94+
- [pymia](https://pymia.readthedocs.io) (low activity)
95+
- [MedicalTorch](https://medicaltorch.readthedocs.io) (abandoned)
96+
- [Eisen](https://github.com/eisen-ai/eisen-core) (abandoned)

‎src/torchio/datasets/__init__.py‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,14 @@
44
from .itk_snap import T1T2
55
from .itk_snap import AorticValve
66
from .itk_snap import BrainTumor
7+
from .ixi import ixi
8+
from .ixi import ixi_tiny
9+
from .medmnist import adrenal_mnist_3d
10+
from .medmnist import fracture_mnist_3d
11+
from .medmnist import nodule_mnist_3d
12+
from .medmnist import organ_mnist_3d
13+
from .medmnist import synapse_mnist_3d
14+
from .medmnist import vessel_mnist_3d
715
from .mni import Colin27
816
from .mni import ICBM2009CNonlinearSymmetric
917
from .mni import Pediatric
@@ -22,4 +30,12 @@
2230
"Sheep",
2331
"Slicer",
2432
"ZonePlate",
33+
"adrenal_mnist_3d",
34+
"fracture_mnist_3d",
35+
"ixi",
36+
"ixi_tiny",
37+
"nodule_mnist_3d",
38+
"organ_mnist_3d",
39+
"synapse_mnist_3d",
40+
"vessel_mnist_3d",
2541
]

‎src/torchio/datasets/ixi.py‎

Lines changed: 184 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,184 @@
1+
"""IXI dataset: ~600 brain MRIs from healthy subjects.
2+
3+
The `Information eXtraction from Images (IXI)
4+
<https://brain-development.org/ixi-dataset/>`_ dataset contains
5+
nearly 600 MR images from normal, healthy subjects.
6+
7+
This data is made available under the Creative Commons CC BY-SA 3.0
8+
license. If you use it, please acknowledge the source.
9+
"""
10+
11+
from __future__ import annotations
12+
13+
import shutil
14+
from collections.abc import Sequence
15+
from pathlib import Path
16+
from tempfile import NamedTemporaryFile
17+
from typing import Any
18+
19+
from ..data.image import LabelMap
20+
from ..data.image import ScalarImage
21+
from ..data.subject import Subject
22+
from ..download import download_and_extract_archive
23+
from ..types import TypePath
24+
25+
26+
def ixi(
27+
root: TypePath,
28+
*,
29+
download: bool = False,
30+
modalities: Sequence[str] = ("T1", "T2"),
31+
) -> list[Subject]:
32+
"""Download and load the full IXI dataset.
33+
34+
Args:
35+
root: Root directory for the dataset.
36+
download: If ``True``, download the data into ``root``.
37+
modalities: Modalities to include. Must be a subset of
38+
``('T1', 'T2', 'PD', 'MRA', 'DTI')``.
39+
40+
Returns:
41+
List of subjects, one per scan with all requested modalities.
42+
43+
Warning:
44+
The dataset is several GB. Downloading may take a while.
45+
"""
46+
root = Path(root)
47+
md5s = _IXI_MD5
48+
for m in modalities:
49+
if m not in md5s:
50+
msg = f'Modality "{m}" must be one of {tuple(md5s.keys())}'
51+
raise ValueError(msg)
52+
if download:
53+
_download_ixi(root, modalities, md5s)
54+
if not all((root / m).is_dir() for m in modalities):
55+
msg = "Dataset not found. Use download=True to download it"
56+
raise RuntimeError(msg)
57+
return _load_ixi_subjects(root, modalities)
58+
59+
60+
def ixi_tiny(
61+
root: TypePath,
62+
*,
63+
download: bool = False,
64+
) -> list[Subject]:
65+
r"""Download and load IXITiny (566 $T_1$ images + segmentations).
66+
67+
All images have shape $83 \times 44 \times 55$. Useful as a
68+
medical image MNIST for quick experiments.
69+
70+
Args:
71+
root: Root directory for the dataset.
72+
download: If ``True``, download the data into ``root``.
73+
74+
Returns:
75+
List of subjects with ``image`` and ``label`` keys.
76+
"""
77+
root = Path(root)
78+
if download:
79+
_download_ixi_tiny(root)
80+
if not root.is_dir():
81+
msg = "Dataset not found. Use download=True to download it"
82+
raise RuntimeError(msg)
83+
return _load_ixi_tiny_subjects(root)
84+
85+
86+
# ---------------------------------------------------------------------------
87+
# Internal helpers
88+
# ---------------------------------------------------------------------------
89+
90+
_IXI_BASE_URL = (
91+
"http://biomedic.doc.ic.ac.uk/brain-development/downloads/IXI/IXI-{modality}.tar"
92+
)
93+
_IXI_MD5: dict[str, str] = {
94+
"T1": "34901a0593b41dd19c1a1f746eac2d58",
95+
"T2": "e3140d78730ecdd32ba92da48c0a9aaa",
96+
"PD": "88ecd9d1fa33cb4a2278183b42ffd749",
97+
"MRA": "29be7d2fee3998f978a55a9bdaf3407e",
98+
"DTI": "636573825b1c8b9e8c78f1877df3ee66",
99+
}
100+
101+
_IXI_TINY_URL = "https://www.dropbox.com/s/ogxjwjxdv5mieah/ixi_tiny.zip?dl=1"
102+
_IXI_TINY_MD5 = "bfb60f4074283d78622760230bfa1f98"
103+
104+
105+
def _download_ixi(
106+
root: Path,
107+
modalities: Sequence[str],
108+
md5s: dict[str, str],
109+
) -> None:
110+
for modality in modalities:
111+
modality_dir = root / modality
112+
if modality_dir.is_dir():
113+
continue
114+
modality_dir.mkdir(exist_ok=True, parents=True)
115+
url = _IXI_BASE_URL.format(modality=modality)
116+
with NamedTemporaryFile(suffix=".tar", delete=False) as f:
117+
download_and_extract_archive(
118+
url,
119+
download_root=modality_dir,
120+
filename=f.name,
121+
md5=md5s[modality],
122+
)
123+
124+
125+
def _load_ixi_subjects(
126+
root: Path,
127+
modalities: Sequence[str],
128+
) -> list[Subject]:
129+
first = modalities[0]
130+
paths = sorted((root / first).glob("*.nii.gz"))
131+
subjects: list[Subject] = []
132+
for filepath in paths:
133+
sid = _subject_id(filepath)
134+
images: dict[str, Any] = {"subject_id": sid}
135+
images[first] = ScalarImage(filepath)
136+
skip = False
137+
for m in modalities[1:]:
138+
matches = sorted((root / m).glob(f"{sid}-{m}.nii.gz"))
139+
if matches:
140+
images[m] = ScalarImage(matches[0])
141+
else:
142+
skip = True
143+
break
144+
if not skip:
145+
subjects.append(Subject(**images))
146+
return subjects
147+
148+
149+
def _download_ixi_tiny(root: Path) -> None:
150+
if root.is_dir():
151+
return
152+
with NamedTemporaryFile(suffix=".zip", delete=False) as f:
153+
download_and_extract_archive(
154+
_IXI_TINY_URL,
155+
download_root=root,
156+
filename=f.name,
157+
md5=_IXI_TINY_MD5,
158+
)
159+
ixi_tiny_dir = root / "ixi_tiny"
160+
(ixi_tiny_dir / "image").rename(root / "image")
161+
(ixi_tiny_dir / "label").rename(root / "label")
162+
shutil.rmtree(ixi_tiny_dir)
163+
164+
165+
def _load_ixi_tiny_subjects(root: Path) -> list[Subject]:
166+
image_paths = sorted((root / "image").glob("*.nii.gz"))
167+
label_paths = sorted((root / "label").glob("*.nii.gz"))
168+
if not (image_paths and label_paths):
169+
msg = f"Images not found. Remove {root} and try again"
170+
raise FileNotFoundError(msg)
171+
subjects: list[Subject] = []
172+
for img_path, lbl_path in zip(image_paths, label_paths, strict=True):
173+
subjects.append(
174+
Subject(
175+
image=ScalarImage(img_path),
176+
label=LabelMap(lbl_path),
177+
subject_id=_subject_id(img_path),
178+
),
179+
)
180+
return subjects
181+
182+
183+
def _subject_id(path: Path) -> str:
184+
return "-".join(path.name.split("-")[:-1])

‎src/torchio/datasets/medmnist.py‎

Lines changed: 113 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,113 @@
1+
"""3D MedMNIST v2 datasets.
2+
3+
Datasets from `MedMNIST v2: A Large-Scale Lightweight Benchmark for
4+
2D and 3D Biomedical Image Classification
5+
<https://arxiv.org/abs/2110.14795>`_.
6+
7+
Check the `MedMNIST website <https://medmnist.com/>`_ for details
8+
and licensing.
9+
"""
10+
11+
from __future__ import annotations
12+
13+
import numpy as np
14+
import torch
15+
16+
from ..data.image import ScalarImage
17+
from ..data.subject import Subject
18+
from ..download import download_url
19+
from ..download import get_torchio_cache_dir
20+
21+
22+
def _load_medmnist(
23+
class_name: str,
24+
split: str,
25+
) -> list[Subject]:
26+
"""Shared loader for all MedMNIST 3D datasets."""
27+
valid = ("train", "training", "val", "validation", "test", "testing")
28+
if split not in valid:
29+
msg = f"Split must be one of {valid}, got {split!r}"
30+
raise ValueError(msg)
31+
match split:
32+
case "training":
33+
split = "train"
34+
case "validation":
35+
split = "val"
36+
case "testing":
37+
split = "test"
38+
39+
filename = f"{class_name}.npz"
40+
url = f"https://zenodo.org/record/5208230/files/{filename}?download=1"
41+
download_root = get_torchio_cache_dir() / "MedMNIST"
42+
download_url(url, download_root, filename=filename)
43+
path = download_root / filename
44+
npz = np.load(path)
45+
images = npz[f"{split}_images"]
46+
labels = npz[f"{split}_labels"]
47+
subjects: list[Subject] = []
48+
for image, label in zip(images, labels, strict=True):
49+
tensor = torch.from_numpy(image[np.newaxis].copy()).float()
50+
scalar = ScalarImage.from_tensor( # ty: ignore[unresolved-attribute]
51+
tensor,
52+
)
53+
subjects.append(
54+
Subject(
55+
image=scalar,
56+
labels=torch.from_numpy(label.copy()),
57+
),
58+
)
59+
return subjects
60+
61+
62+
def organ_mnist_3d(split: str = "train") -> list[Subject]:
63+
"""3D organ segmentation dataset.
64+
65+
Args:
66+
split: ``'train'``, ``'val'``, or ``'test'``.
67+
"""
68+
return _load_medmnist("organmnist3d", split)
69+
70+
71+
def nodule_mnist_3d(split: str = "train") -> list[Subject]:
72+
"""3D lung nodule dataset.
73+
74+
Args:
75+
split: ``'train'``, ``'val'``, or ``'test'``.
76+
"""
77+
return _load_medmnist("nodulemnist3d", split)
78+
79+
80+
def adrenal_mnist_3d(split: str = "train") -> list[Subject]:
81+
"""3D adrenal gland dataset.
82+
83+
Args:
84+
split: ``'train'``, ``'val'``, or ``'test'``.
85+
"""
86+
return _load_medmnist("adrenalmnist3d", split)
87+
88+
89+
def fracture_mnist_3d(split: str = "train") -> list[Subject]:
90+
"""3D bone fracture dataset.
91+
92+
Args:
93+
split: ``'train'``, ``'val'``, or ``'test'``.
94+
"""
95+
return _load_medmnist("fracturemnist3d", split)
96+
97+
98+
def vessel_mnist_3d(split: str = "train") -> list[Subject]:
99+
"""3D vessel dataset.
100+
101+
Args:
102+
split: ``'train'``, ``'val'``, or ``'test'``.
103+
"""
104+
return _load_medmnist("vesselmnist3d", split)
105+
106+
107+
def synapse_mnist_3d(split: str = "train") -> list[Subject]:
108+
"""3D synapse dataset.
109+
110+
Args:
111+
split: ``'train'``, ``'val'``, or ``'test'``.
112+
"""
113+
return _load_medmnist("synapsemnist3d", split)

‎src/torchio/datasets/mni/colin.py‎

Lines changed: 4 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -76,16 +76,10 @@ def __init__(self, version: int = 1998) -> None:
7676
compress(path)
7777
path.unlink()
7878

79-
try:
80-
subject_kwargs = self._get_subject_kwargs(
81-
self.download_root,
82-
extension=".nii.gz",
83-
)
84-
except FileNotFoundError: # for backward compatibility
85-
subject_kwargs = self._get_subject_kwargs(
86-
self.download_root,
87-
extension=".nii",
88-
)
79+
subject_kwargs = self._get_subject_kwargs(
80+
self.download_root,
81+
extension=".nii.gz",
82+
)
8983
super().__init__(**subject_kwargs)
9084

9185
def _get_subject_kwargs(self, download_root, extension):

‎src/torchio/datasets/mni/pediatric.py‎

Lines changed: 2 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -67,10 +67,7 @@ def __init__(
6767
compress(path)
6868
path.unlink()
6969

70-
try:
71-
subject_kwargs = self._get_subject_kwargs(".nii.gz")
72-
except FileNotFoundError: # for backward compatibility
73-
subject_kwargs = self._get_subject_kwargs(".nii")
70+
subject_kwargs = self._get_subject_kwargs(".nii.gz")
7471
super().__init__(**subject_kwargs)
7572

7673
def _get_subject_kwargs(self, extension: str) -> dict:
@@ -79,7 +76,5 @@ def _get_subject_kwargs(self, extension: str) -> dict:
7976
"t1": ScalarImage(root / f"nihpd_{self.file_id}_t1w{extension}"),
8077
"t2": ScalarImage(root / f"nihpd_{self.file_id}_t2w{extension}"),
8178
"pd": ScalarImage(root / f"nihpd_{self.file_id}_pdw{extension}"),
82-
"mask": LabelMap(
83-
root / f"nihpd_{self.file_id}_mask{extension}",
84-
),
79+
"mask": LabelMap(root / f"nihpd_{self.file_id}_mask{extension}"),
8580
}

‎src/torchio/datasets/mni/sheep.py‎

Lines changed: 1 addition & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -36,7 +36,4 @@ def __init__(self) -> None:
3636
if path == t1_nii_path:
3737
compress(t1_nii_path, t1_niigz_path)
3838
path.unlink()
39-
try:
40-
super().__init__(t1=ScalarImage(t1_niigz_path))
41-
except FileNotFoundError: # for backward compatibility
42-
super().__init__(t1=ScalarImage(t1_nii_path))
39+
super().__init__(t1=ScalarImage(t1_niigz_path))

0 commit comments

Comments
 (0)