torch_em.data.datasets.medical.nasalseg

  1"""The NasalSeg dataset contains annotations for nasal cavity and paranasal sinus segmentation in 3D CT scans.
  2
  3The dataset consists of 130 CT scans (NRRD, 'images/<case>_img.nrrd' and 'labels/<case>_seg.nrrd') with voxel-wise
  4annotations of five structures. The label ids are described in `LABEL_IDS`. The record does not document the id
  5order; it was determined from the label positions in the scans (the images use the left-posterior-superior
  6convention): 1 and 2 are the lateral pair (maxillary sinuses), 3 and 4 the medial pair (nasal cavities) and 5 is
  7the posterior midline structure (nasopharynx), with the lower id of each pair on the right side of the patient.
  8The scans are converted to hdf5 files (keys 'raw' and 'labels') by this module.
  9
 10The dataset is located at https://doi.org/10.5281/zenodo.13893419, released under a CC-BY-4.0 license.
 11Please cite the corresponding publication if you use this dataset for your research.
 12"""
 13
 14import os
 15from glob import glob
 16from natsort import natsorted
 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/api/records/13893419/files/NasalSeg.zip/content"
 27CHECKSUM = "60c6facf843685802c39e4adff4a05c081c1c4b6175c9cb573745c55abb0fa6a"
 28
 29LABEL_IDS = {
 30    "right maxillary sinus": 1,
 31    "left maxillary sinus": 2,
 32    "right nasal cavity": 3,
 33    "left nasal cavity": 4,
 34    "nasopharynx": 5,
 35}
 36
 37
 38def _convert_case(image_path, label_path, out_path):
 39    import h5py
 40    import SimpleITK as sitk
 41
 42    if os.path.exists(out_path):
 43        return
 44
 45    raw = sitk.GetArrayFromImage(sitk.ReadImage(image_path))
 46    labels = sitk.GetArrayFromImage(sitk.ReadImage(label_path)).astype("uint8")
 47
 48    tmp_path = f"{out_path}.{os.getpid()}.incomplete"
 49    with h5py.File(tmp_path, "w") as f:
 50        f.create_dataset("raw", data=raw, compression="gzip")
 51        f.create_dataset("labels", data=labels, compression="gzip")
 52    os.replace(tmp_path, out_path)
 53
 54
 55def get_nasalseg_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 56    """Download the NasalSeg dataset.
 57
 58    Args:
 59        path: Filepath to a folder where the data is downloaded for further processing.
 60        download: Whether to download the data if it is not present.
 61
 62    Returns:
 63        Filepath where the converted data is stored.
 64    """
 65    image_dir = os.path.join(path, "images")
 66    label_dir = os.path.join(path, "labels")
 67    if not os.path.exists(image_dir) or not os.path.exists(label_dir):
 68        os.makedirs(path, exist_ok=True)
 69        zip_path = os.path.join(path, "NasalSeg.zip")
 70        util.download_source(path=zip_path, url=URL, download=download, checksum=CHECKSUM)
 71        util.unzip(zip_path=zip_path, dst=path, remove=False)
 72
 73    assert os.path.exists(image_dir) and os.path.exists(label_dir), \
 74        f"The extraction of the NasalSeg archive did not create the expected folders in '{path}'."
 75
 76    converted_dir = os.path.join(path, "converted")
 77    os.makedirs(converted_dir, exist_ok=True)
 78    for image_path in natsorted(glob(os.path.join(image_dir, "*_img.nrrd"))):
 79        case_id = os.path.basename(image_path)[:-len("_img.nrrd")]
 80        label_path = os.path.join(label_dir, f"{case_id}_seg.nrrd")
 81        _convert_case(image_path, label_path, os.path.join(converted_dir, f"{case_id}.h5"))
 82
 83    return converted_dir
 84
 85
 86def get_nasalseg_paths(path: Union[os.PathLike, str], download: bool = False) -> List[str]:
 87    """Get paths to the NasalSeg data.
 88
 89    Args:
 90        path: Filepath to a folder where the data is downloaded for further processing.
 91        download: Whether to download the data if it is not present.
 92
 93    Returns:
 94        List of filepaths for the hdf5 files, which contain the image data ('raw') and the label data ('labels').
 95    """
 96    converted_dir = get_nasalseg_data(path, download)
 97    volume_paths = natsorted(glob(os.path.join(converted_dir, "*.h5")))
 98    assert len(volume_paths) > 0
 99    return volume_paths
100
101
102def get_nasalseg_dataset(
103    path: Union[os.PathLike, str],
104    patch_shape: Tuple[int, int, int],
105    resize_inputs: bool = False,
106    download: bool = False,
107    **kwargs
108) -> Dataset:
109    """Get the NasalSeg dataset for nasal cavity and paranasal sinus segmentation.
110
111    Args:
112        path: Filepath to a folder where the data is downloaded for further processing.
113        patch_shape: The patch shape to use for training.
114        resize_inputs: Whether to resize the inputs to the patch shape.
115        download: Whether to download the data if it is not present.
116        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
117
118    Returns:
119        The segmentation dataset.
120    """
121    volume_paths = get_nasalseg_paths(path, download)
122
123    if resize_inputs:
124        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
125        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
126            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
127        )
128
129    return torch_em.default_segmentation_dataset(
130        raw_paths=volume_paths,
131        raw_key="raw",
132        label_paths=volume_paths,
133        label_key="labels",
134        is_seg_dataset=True,
135        patch_shape=patch_shape,
136        ndim=3,
137        **kwargs
138    )
139
140
141def get_nasalseg_loader(
142    path: Union[os.PathLike, str],
143    batch_size: int,
144    patch_shape: Tuple[int, int, int],
145    resize_inputs: bool = False,
146    download: bool = False,
147    **kwargs
148) -> DataLoader:
149    """Get the NasalSeg dataloader for nasal cavity and paranasal sinus segmentation.
150
151    Args:
152        path: Filepath to a folder where the data is downloaded for further processing.
153        batch_size: The batch size for training.
154        patch_shape: The patch shape to use for training.
155        resize_inputs: Whether to resize the inputs to the patch shape.
156        download: Whether to download the data if it is not present.
157        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
158
159    Returns:
160        The DataLoader.
161    """
162    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
163    dataset = get_nasalseg_dataset(path, patch_shape, resize_inputs, download, **ds_kwargs)
164    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)