torch_em.data.datasets.medical.hntsmrg

  1"""The HNTS-MRG dataset contains annotations for head and neck tumor segmentation in T2-weighted MRI.
  2
  3The dataset is the training data of the HNTS-MRG 2024 challenge (https://hntsmrg24.grand-challenge.org).
  4It consists of 150 patients, each with a pre-radiotherapy (pre-RT) and a mid-radiotherapy (mid-RT) T2-weighted MRI
  5and the corresponding tumor segmentation. The labels are: 0 = background, 1 = primary gross tumor volume (GTVp),
  62 = metastatic lymph nodes (GTVn). NOTE: The mid-RT masks of the patients 21, 25, 29 and 42 are empty
  7(complete response to the therapy). For the mid-RT timepoint the dataset additionally provides the pre-RT image and
  8mask registered to the mid-RT image space ('*_preRT_T2_registered.nii.gz', '*_preRT_mask_registered.nii.gz').
  9
 10The dataset is located at https://doi.org/10.5281/zenodo.11199559.
 11
 12This dataset is from the publication https://doi.org/10.1007/978-3-031-83274-1_1.
 13Please cite it if you use this dataset in your research.
 14"""
 15
 16import os
 17from glob import glob
 18from natsort import natsorted
 19from typing import Union, Tuple, Literal, List
 20
 21from torch.utils.data import Dataset, DataLoader
 22
 23import torch_em
 24
 25from .. import util
 26
 27
 28URL = "https://zenodo.org/records/11199559/files/HNTSMRG24_train.zip"
 29CHECKSUM = "675e3c56509b0b5bcaad1291990dfa4a8ef054bb468a71b6cd08189a5e67bfb6"
 30
 31
 32def get_hntsmrg_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 33    """Download the HNTS-MRG dataset.
 34
 35    Args:
 36        path: Filepath to a folder where the data is downloaded for further processing.
 37        download: Whether to download the data if it is not present.
 38
 39    Returns:
 40        Filepath where the data is downloaded.
 41    """
 42    data_dir = os.path.join(path, "HNTSMRG24_train")
 43    if os.path.exists(data_dir):
 44        return data_dir
 45
 46    os.makedirs(path, exist_ok=True)
 47
 48    zip_path = os.path.join(path, "HNTSMRG24_train.zip")
 49    util.download_source(path=zip_path, url=URL, download=download, checksum=CHECKSUM)
 50    util.unzip(zip_path=zip_path, dst=path)
 51
 52    return data_dir
 53
 54
 55def get_hntsmrg_paths(
 56    path: Union[os.PathLike, str], timepoint: Literal['pre', 'mid'] = "pre", download: bool = False
 57) -> Tuple[List[str], List[str]]:
 58    """Get paths to the HNTS-MRG data.
 59
 60    Args:
 61        path: Filepath to a folder where the data is downloaded for further processing.
 62        timepoint: The timepoint of the scans. Either 'pre' (pre-radiotherapy) or 'mid' (mid-radiotherapy).
 63        download: Whether to download the data if it is not present.
 64
 65    Returns:
 66        List of filepaths for the image data.
 67        List of filepaths for the label data.
 68    """
 69    data_dir = get_hntsmrg_data(path, download)
 70
 71    if timepoint not in ("pre", "mid"):
 72        raise ValueError(f"'{timepoint}' is not a valid timepoint. Choose either 'pre' or 'mid'.")
 73
 74    tp = f"{timepoint}RT"
 75    raw_paths = natsorted(glob(os.path.join(data_dir, "*", tp, f"*_{tp}_T2.nii.gz")))
 76    label_paths = [p.replace(f"_{tp}_T2.nii.gz", f"_{tp}_mask.nii.gz") for p in raw_paths]
 77    assert len(raw_paths) > 0 and all(os.path.exists(p) for p in label_paths)
 78
 79    return raw_paths, label_paths
 80
 81
 82def get_hntsmrg_dataset(
 83    path: Union[os.PathLike, str],
 84    patch_shape: Tuple[int, ...],
 85    timepoint: Literal['pre', 'mid'] = "pre",
 86    resize_inputs: bool = False,
 87    download: bool = False,
 88    **kwargs
 89) -> Dataset:
 90    """Get the HNTS-MRG dataset for head and neck tumor segmentation in MRI.
 91
 92    Args:
 93        path: Filepath to a folder where the data is downloaded for further processing.
 94        patch_shape: The patch shape to use for training.
 95        timepoint: The timepoint of the scans. Either 'pre' (pre-radiotherapy) or 'mid' (mid-radiotherapy).
 96        resize_inputs: Whether to resize inputs to the desired patch shape.
 97        download: Whether to download the data if it is not present.
 98        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
 99
100    Returns:
101        The segmentation dataset.
102    """
103    raw_paths, label_paths = get_hntsmrg_paths(path, timepoint, download)
104
105    if resize_inputs:
106        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
107        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
108            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
109        )
110
111    return torch_em.default_segmentation_dataset(
112        raw_paths=raw_paths,
113        raw_key="data",
114        label_paths=label_paths,
115        label_key="data",
116        patch_shape=patch_shape,
117        is_seg_dataset=True,
118        **kwargs
119    )
120
121
122def get_hntsmrg_loader(
123    path: Union[os.PathLike, str],
124    batch_size: int,
125    patch_shape: Tuple[int, ...],
126    timepoint: Literal['pre', 'mid'] = "pre",
127    resize_inputs: bool = False,
128    download: bool = False,
129    **kwargs
130) -> DataLoader:
131    """Get the HNTS-MRG dataloader for head and neck tumor segmentation in MRI.
132
133    Args:
134        path: Filepath to a folder where the data is downloaded for further processing.
135        batch_size: The batch size for training.
136        patch_shape: The patch shape to use for training.
137        timepoint: The timepoint of the scans. Either 'pre' (pre-radiotherapy) or 'mid' (mid-radiotherapy).
138        resize_inputs: Whether to resize inputs to the desired patch shape.
139        download: Whether to download the data if it is not present.
140        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
141
142    Returns:
143        The DataLoader.
144    """
145    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
146    dataset = get_hntsmrg_dataset(path, patch_shape, timepoint, resize_inputs, download, **ds_kwargs)
147    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)