torch_em.data.datasets.medical.ibd_mre

  1"""The IBD-MRE dataset contains annotations for bowel segment segmentation in
  2magnetic resonance enterography (MRE) scans of patients with inflammatory bowel disease (IBD).
  3
  4The dataset contains coronal HASTE (half-Fourier acquisition single-shot turbo spin-echo) MRE
  5sequences from 114 IBD patients, with fine pixel-level annotations for ten bowel segments,
  6labeled by experienced radiologists. The semantic label ids are:
  71: stomach, 2: duodenum, 3: small intestine, 4: appendix, 5: cecum, 6: ascending colon,
  87: transverse colon, 8: descending colon, 9: sigmoid colon, 10: rectum.
  9
 10The dataset is located at https://doi.org/10.5281/zenodo.13839321.
 11This dataset is from the publication https://doi.org/10.1038/s41597-025-04760-z.
 12Please cite it if you use this dataset for your research.
 13"""
 14
 15import os
 16from glob import glob
 17from typing import Union, Tuple, List
 18
 19from torch.utils.data import Dataset, DataLoader
 20
 21import torch_em
 22
 23from .. import util
 24
 25
 26URL = "https://zenodo.org/records/13839321/files/A%20comprehensive%20dataset.zip"
 27CHECKSUM = "d88541e64f33629b8390addd286f7dea88e665d29c2cec21d7594b7c64c97393"
 28
 29
 30def get_ibd_mre_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 31    """Download the IBD-MRE dataset.
 32
 33    Args:
 34        path: Filepath to a folder where the data is downloaded for further processing.
 35        download: Whether to download the data if it is not present.
 36
 37    Returns:
 38        Filepath where the data is downloaded.
 39    """
 40    data_dir = os.path.join(path, "data")
 41    if os.path.exists(data_dir):
 42        return data_dir
 43
 44    os.makedirs(path, exist_ok=True)
 45
 46    zip_path = os.path.join(path, "A comprehensive dataset.zip")
 47    util.download_source(path=zip_path, url=URL, download=download, checksum=CHECKSUM)
 48    util.unzip(zip_path=zip_path, dst=data_dir)
 49
 50    return data_dir
 51
 52
 53def get_ibd_mre_paths(path: Union[os.PathLike, str], download: bool = False) -> Tuple[List[str], List[str]]:
 54    """Get paths to the IBD-MRE data.
 55
 56    Args:
 57        path: Filepath to a folder where the data is downloaded for further processing.
 58        download: Whether to download the data if it is not present.
 59
 60    Returns:
 61        List of filepaths for the image data.
 62        List of filepaths for the label data.
 63    """
 64    data_dir = get_ibd_mre_data(path=path, download=download)
 65
 66    image_paths = sorted(
 67        glob(os.path.join(data_dir, "*_data.nii.gz")), key=lambda p: int(os.path.basename(p).split("_")[0])
 68    )
 69    gt_paths = sorted(
 70        glob(os.path.join(data_dir, "*_label.nii.gz")), key=lambda p: int(os.path.basename(p).split("_")[0])
 71    )
 72
 73    return image_paths, gt_paths
 74
 75
 76def get_ibd_mre_dataset(
 77    path: Union[os.PathLike, str], patch_shape: Tuple[int, ...], download: bool = False, **kwargs
 78) -> Dataset:
 79    """Get the IBD-MRE dataset for segmentation of bowel segments in MRE scans.
 80
 81    Args:
 82        path: Filepath to a folder where the data is downloaded for further processing.
 83        patch_shape: The patch shape to use for training.
 84        download: Whether to download the data if it is not present.
 85        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
 86
 87    Returns:
 88        The segmentation dataset.
 89    """
 90    image_paths, gt_paths = get_ibd_mre_paths(path, download)
 91
 92    return torch_em.default_segmentation_dataset(
 93        raw_paths=image_paths,
 94        raw_key="data",
 95        label_paths=gt_paths,
 96        label_key="data",
 97        patch_shape=patch_shape,
 98        is_seg_dataset=True,
 99        **kwargs
100    )
101
102
103def get_ibd_mre_loader(
104    path: Union[os.PathLike, str], batch_size: int, patch_shape: Tuple[int, ...], download: bool = False, **kwargs
105) -> DataLoader:
106    """Get the IBD-MRE dataloader for segmentation of bowel segments in MRE scans.
107
108    Args:
109        path: Filepath to a folder where the data is downloaded for further processing.
110        batch_size: The batch size for training.
111        patch_shape: The patch shape to use for training.
112        download: Whether to download the data if it is not present.
113        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
114
115    Returns:
116        The DataLoader.
117    """
118    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
119    dataset = get_ibd_mre_dataset(path, patch_shape, download, **ds_kwargs)
120    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)