torch_em.data.datasets.medical.mama_mia

  1"""The MAMA-MIA dataset contains annotations for primary breast tumor segmentation in
  2dynamic contrast-enhanced (DCE) MRI.
  3
  4The dataset pools 1506 pre-treatment breast DCE-MRI of four public collections of The Cancer Imaging Archive
  5(DUKE, ISPY1, ISPY2 and NACT), for which the primary tumor was segmented by sixteen breast radiologists. The
  6segmentation is a binary mask of the primary tumor (which includes the non-mass enhancement areas), see
  7`LABEL_IDS`. The collection a case belongs to is encoded in its name and can be selected with the 'cohort'
  8argument.
  9
 10NOTE: The official release at https://www.synapse.org/Synapse:syn60868042 requires a Synapse (or Health-RI XNAT)
 11account, so it cannot be downloaded automatically. This module uses the open redistribution at
 12https://huggingface.co/datasets/YongchengYAO/MAMA-MIA-Lite (CC BY-NC 4.0), which contains all 1506 cases with
 13their expert segmentations, but only the first post-contrast DCE phase (the phase the masks were drawn on) of
 14each case. The other DCE phases and the preliminary automatic segmentations of the official release are not part
 15of this redistribution, so all masks provided here are the expert segmentations. The volumes of the
 16redistribution were reoriented to RAS+ and the masks were cast to uint16.
 17
 18The volumes are stored as nifti files with the slice axis last, but they are loaded with the slice axis first
 19(torch_em reverses the nifti axis order), so 2d patches are extracted along the axial axis.
 20
 21This dataset is from the publication https://doi.org/10.1038/s41597-025-04707-4.
 22Please cite it if you use this dataset in your research.
 23"""
 24
 25import os
 26from glob import glob
 27from natsort import natsorted
 28from typing import Union, Tuple, List, Literal, Optional
 29
 30from torch.utils.data import Dataset, DataLoader
 31
 32import torch_em
 33
 34from .. import util
 35
 36
 37URLS = {
 38    "part1": "https://huggingface.co/datasets/YongchengYAO/MAMA-MIA-Lite/resolve/main/data-part001.zip",
 39    "part2": "https://huggingface.co/datasets/YongchengYAO/MAMA-MIA-Lite/resolve/main/data-part002.zip",
 40}
 41
 42CHECKSUMS = {
 43    "part1": "8cbd3d0165a446793b9f342ef28dc27be4a7a81e4933f69515b5287fe6b0b89c",
 44    "part2": "22c87ab82574279c235f7863e9d2de0e8baa031ca998278b6a0a30722dabffee",
 45}
 46
 47LABEL_IDS = {"background": 0, "primary_tumor": 1}
 48
 49COHORTS = ["DUKE", "ISPY1", "ISPY2", "NACT"]
 50
 51N_VOLUMES = 1506
 52
 53
 54def get_mama_mia_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 55    """Download the MAMA-MIA dataset.
 56
 57    Args:
 58        path: Filepath to a folder where the data is downloaded for further processing.
 59        download: Whether to download the data if it is not present.
 60
 61    Returns:
 62        Filepath where the data is stored.
 63    """
 64    if len(glob(os.path.join(path, "Images", "*.nii.gz"))) == N_VOLUMES:
 65        return path
 66
 67    os.makedirs(path, exist_ok=True)
 68
 69    for part, url in URLS.items():
 70        zip_path = os.path.join(path, f"{part}.zip")
 71        util.download_source(path=zip_path, url=url, download=download, checksum=CHECKSUMS[part])
 72        util.unzip(zip_path=zip_path, dst=path)
 73
 74    return path
 75
 76
 77def get_mama_mia_paths(
 78    path: Union[os.PathLike, str],
 79    cohort: Optional[Literal["DUKE", "ISPY1", "ISPY2", "NACT"]] = None,
 80    download: bool = False,
 81) -> Tuple[List[str], List[str]]:
 82    """Get paths to the MAMA-MIA data.
 83
 84    Args:
 85        path: Filepath to a folder where the data is downloaded for further processing.
 86        cohort: The choice of source collection. By default all four collections are used.
 87        download: Whether to download the data if it is not present.
 88
 89    Returns:
 90        List of filepaths for the image data.
 91        List of filepaths for the label data.
 92    """
 93    if cohort is not None and cohort not in COHORTS:
 94        raise ValueError(f"'{cohort}' is not a valid cohort. Please choose one of {COHORTS}.")
 95
 96    data_dir = get_mama_mia_data(path, download)
 97
 98    prefix = "*" if cohort is None else f"{cohort}_*"
 99    raw_paths = natsorted(glob(os.path.join(data_dir, "Images", f"{prefix}.nii.gz")))
100    label_paths = [os.path.join(data_dir, "Masks", os.path.basename(p)) for p in raw_paths]
101    assert len(raw_paths) > 0 and all(os.path.exists(p) for p in label_paths)
102
103    return raw_paths, label_paths
104
105
106def get_mama_mia_dataset(
107    path: Union[os.PathLike, str],
108    patch_shape: Tuple[int, ...],
109    cohort: Optional[Literal["DUKE", "ISPY1", "ISPY2", "NACT"]] = None,
110    resize_inputs: bool = False,
111    download: bool = False,
112    **kwargs
113) -> Dataset:
114    """Get the MAMA-MIA dataset for breast tumor segmentation.
115
116    Args:
117        path: Filepath to a folder where the data is downloaded for further processing.
118        patch_shape: The patch shape to use for training.
119        cohort: The choice of source collection. By default all four collections are used.
120        resize_inputs: Whether to resize inputs to the desired patch shape.
121        download: Whether to download the data if it is not present.
122        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
123
124    Returns:
125        The segmentation dataset.
126    """
127    raw_paths, label_paths = get_mama_mia_paths(path, cohort, download)
128
129    if resize_inputs:
130        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
131        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
132            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
133        )
134
135    return torch_em.default_segmentation_dataset(
136        raw_paths=raw_paths,
137        raw_key="data",
138        label_paths=label_paths,
139        label_key="data",
140        patch_shape=patch_shape,
141        is_seg_dataset=True,
142        **kwargs
143    )
144
145
146def get_mama_mia_loader(
147    path: Union[os.PathLike, str],
148    batch_size: int,
149    patch_shape: Tuple[int, ...],
150    cohort: Optional[Literal["DUKE", "ISPY1", "ISPY2", "NACT"]] = None,
151    resize_inputs: bool = False,
152    download: bool = False,
153    **kwargs
154) -> DataLoader:
155    """Get the MAMA-MIA dataloader for breast tumor segmentation.
156
157    Args:
158        path: Filepath to a folder where the data is downloaded for further processing.
159        batch_size: The batch size for training.
160        patch_shape: The patch shape to use for training.
161        cohort: The choice of source collection. By default all four collections are used.
162        resize_inputs: Whether to resize inputs to the desired patch shape.
163        download: Whether to download the data if it is not present.
164        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
165
166    Returns:
167        The DataLoader.
168    """
169    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
170    dataset = get_mama_mia_dataset(path, patch_shape, cohort, resize_inputs, download, **ds_kwargs)
171    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)