torch_em.data.datasets.medical.tom500

  1"""The TOM500 dataset contains annotations for multi-organ orbital segmentation
  2in T2-weighted orbital MRI for thyroid eye disease.
  3
  4The dataset consists of 500 patients with thyroid eye disease, each with a coronal T2-weighted MRI scan
  5(512 x 512 x 20 voxels) and corresponding segmentations of nine orbital structures: the optic nerve,
  6orbital fat, lacrimal gland, eyeball and five extraocular muscles (superior rectus and levator palpebrae
  7superioris complex, inferior rectus, medial rectus, lateral rectus, superior oblique). The annotations
  8were created by three junior annotators and reviewed by an expert radiologist. The data is split into
  9400 training and 100 validation scans.
 10
 11The dataset is located at https://doi.org/10.6084/m9.figshare.27133389 and is distributed under the
 12CC0 1.0 license.
 13
 14This dataset is from the publication https://doi.org/10.1038/s41597-025-04427-9.
 15Please cite it if you use this dataset in your research.
 16"""
 17
 18import os
 19from glob import glob
 20from natsort import natsorted
 21from typing import Union, Tuple, Literal, List
 22
 23from torch.utils.data import Dataset, DataLoader
 24
 25import torch_em
 26
 27from .. import util
 28
 29
 30URL = "https://ndownloader.figshare.com/files/49499655"
 31CHECKSUM = "80bbd9934584573aabef3525ab1736d1c3ed7d9dcc9e8cdff28b3d22627e089f"
 32
 33
 34def get_tom500_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 35    """Download the TOM500 dataset.
 36
 37    Args:
 38        path: Filepath to a folder where the data is downloaded for further processing.
 39        download: Whether to download the data if it is not present.
 40
 41    Returns:
 42        Filepath where the data is downloaded.
 43    """
 44    data_dir = os.path.join(path, "TOM500")
 45    if os.path.exists(data_dir):
 46        return data_dir
 47
 48    os.makedirs(path, exist_ok=True)
 49
 50    zip_path = os.path.join(path, "TOM500.zip")
 51    util.download_source(path=zip_path, url=URL, download=download, checksum=CHECKSUM)
 52    util.unzip(zip_path=zip_path, dst=path)
 53
 54    return data_dir
 55
 56
 57def get_tom500_paths(
 58    path: Union[os.PathLike, str], split: Literal['train', 'val'], download: bool = False
 59) -> Tuple[List[str], List[str]]:
 60    """Get paths to the TOM500 data.
 61
 62    Args:
 63        path: Filepath to a folder where the data is downloaded for further processing.
 64        split: The choice of data split. Either 'train' or 'val'.
 65        download: Whether to download the data if it is not present.
 66
 67    Returns:
 68        List of filepaths for the image data.
 69        List of filepaths for the label data.
 70    """
 71    data_dir = get_tom500_data(path, download)
 72
 73    if split not in ("train", "val"):
 74        raise ValueError(f"'{split}' is not a valid split. Choose either 'train' or 'val'.")
 75
 76    raw_paths = natsorted(glob(os.path.join(data_dir, split, "image", "*.nii.gz")))
 77    label_paths = natsorted(glob(os.path.join(data_dir, split, "label", "*.nii.gz")))
 78    assert len(raw_paths) > 0 and len(raw_paths) == len(label_paths)
 79
 80    return raw_paths, label_paths
 81
 82
 83def get_tom500_dataset(
 84    path: Union[os.PathLike, str],
 85    patch_shape: Tuple[int, ...],
 86    split: Literal['train', 'val'],
 87    resize_inputs: bool = False,
 88    download: bool = False,
 89    **kwargs
 90) -> Dataset:
 91    """Get the TOM500 dataset for multi-organ orbital segmentation in MRI.
 92
 93    Args:
 94        path: Filepath to a folder where the data is downloaded for further processing.
 95        patch_shape: The patch shape to use for training.
 96        split: The choice of data split. Either 'train' or 'val'.
 97        resize_inputs: Whether to resize inputs to the desired patch shape.
 98        download: Whether to download the data if it is not present.
 99        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
100
101    Returns:
102        The segmentation dataset.
103    """
104    raw_paths, label_paths = get_tom500_paths(path, split, download)
105
106    if resize_inputs:
107        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
108        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
109            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
110        )
111
112    return torch_em.default_segmentation_dataset(
113        raw_paths=raw_paths,
114        raw_key="data",
115        label_paths=label_paths,
116        label_key="data",
117        patch_shape=patch_shape,
118        is_seg_dataset=True,
119        **kwargs
120    )
121
122
123def get_tom500_loader(
124    path: Union[os.PathLike, str],
125    batch_size: int,
126    patch_shape: Tuple[int, ...],
127    split: Literal['train', 'val'],
128    resize_inputs: bool = False,
129    download: bool = False,
130    **kwargs
131) -> DataLoader:
132    """Get the TOM500 dataloader for multi-organ orbital segmentation in MRI.
133
134    Args:
135        path: Filepath to a folder where the data is downloaded for further processing.
136        batch_size: The batch size for training.
137        patch_shape: The patch shape to use for training.
138        split: The choice of data split. Either 'train' or 'val'.
139        resize_inputs: Whether to resize inputs to the desired patch shape.
140        download: Whether to download the data if it is not present.
141        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
142
143    Returns:
144        The DataLoader.
145    """
146    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
147    dataset = get_tom500_dataset(path, patch_shape, split, resize_inputs, download, **ds_kwargs)
148    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)