torch_em.data.datasets.medical.hva_ct

  1"""The HVA-CT dataset contains annotations for intrahepatic veins and the liver in CT scans.
  2
  3The annotations re-annotate the portal and hepatic venous systems of the 61 thin-slice scans of the
  4Medical Segmentation Decathlon task 8 (hepatic vessel), so the images are downloaded from there. Two
  5annotation sets are provided: 'vessels' with 1: portal vein and 2: hepatic vein, and 'liver' with a
  6binary liver mask. See also `CLASS_IDS`.
  7
  8NOTE: The ids of the two venous systems are the reverse of the order that the file names suggest. They
  9were assigned from the per-case voxel counts that the release ships in its metadata, which match id 1 to
 10the portal and id 2 to the hepatic vein for all 61 scans.
 11
 12The dataset is located at https://doi.org/10.5281/zenodo.19850108 and is distributed under the
 13CC BY-SA 4.0 license.
 14Please cite it, and the Medical Segmentation Decathlon publication
 15https://doi.org/10.1038/s41467-022-30695-9, if you use this dataset in your research.
 16"""
 17
 18import os
 19from glob import glob
 20from natsort import natsorted
 21from typing import Union, Optional, Tuple, Literal, List
 22
 23from torch.utils.data import Dataset, DataLoader
 24
 25import torch_em
 26
 27from .msd import get_msd_data
 28from .. import util
 29
 30
 31URLS = {
 32    "vessels": "https://zenodo.org/records/19850108/files/hp_masks.zip?download=1",
 33    "liver": "https://zenodo.org/records/19850108/files/liver_masks.zip?download=1",
 34}
 35
 36CHECKSUMS = {
 37    "vessels": "539d5a83c5b9f6f8d890727923d0666e53b8b585f02651a6f250689055fbb4b1",
 38    "liver": "25c40378fde4f1a1a52518fec7d930a4588b904136893ada69004ffe231efe96",
 39}
 40
 41ANNOTATIONS = {
 42    "vessels": ("hp_masks", "hepatic_portalvessel_"),
 43    "liver": ("liver_masks", "liver_"),
 44}
 45"""Mapping from the annotation choice to its folder and the prefix of its files."""
 46
 47CLASS_NAMES = ["portal_vein", "hepatic_vein"]
 48"""The venous systems of the 'vessels' annotations. The label id of a system is its 1-based index."""
 49
 50CLASS_IDS = {name: i + 1 for i, name in enumerate(CLASS_NAMES)}
 51"""Mapping from the venous system to its label id."""
 52
 53
 54def _get_image_paths(msd_dir):
 55    image_paths = {}
 56    for split in ["imagesTr", "imagesTs"]:
 57        for path in glob(os.path.join(msd_dir, "Task08_HepaticVessel", split, "*.nii.gz")):
 58            fname = os.path.basename(path)
 59            # The MSD archives carry macOS resource fork files next to the actual volumes.
 60            if fname.startswith("._"):
 61                continue
 62            image_paths[fname[len("hepaticvessel_"):-len(".nii.gz")]] = path
 63    return image_paths
 64
 65
 66def get_hva_ct_data(
 67    path: Union[os.PathLike, str],
 68    annotation: Literal["vessels", "liver"] = "vessels",
 69    msd_path: Optional[Union[os.PathLike, str]] = None,
 70    download: bool = False,
 71) -> Tuple[str, str]:
 72    """Download the HVA-CT dataset.
 73
 74    Args:
 75        path: Filepath to a folder where the data is downloaded for further processing.
 76        annotation: The choice of annotations. Either 'vessels' or 'liver'.
 77        msd_path: Filepath to an existing download of the Medical Segmentation Decathlon. The scans are
 78            downloaded to `path` if it is not given.
 79        download: Whether to download the data if it is not present.
 80
 81    Returns:
 82        Filepath where the annotations are downloaded.
 83        Filepath where the scans are downloaded.
 84    """
 85    if annotation not in ANNOTATIONS:
 86        raise ValueError(f"'{annotation}' is not a valid annotation. Choose from {list(ANNOTATIONS.keys())}.")
 87
 88    folder, _ = ANNOTATIONS[annotation]
 89    label_dir = os.path.join(path, folder)
 90    if not os.path.exists(label_dir):
 91        os.makedirs(path, exist_ok=True)
 92        zip_path = os.path.join(path, f"{folder}.zip")
 93        util.download_source(
 94            path=zip_path, url=URLS[annotation], download=download, checksum=CHECKSUMS[annotation]
 95        )
 96        util.unzip(zip_path=zip_path, dst=path, remove=False)
 97
 98    # The scans are the hepatic vessel scans of the Medical Segmentation Decathlon.
 99    msd_dir = get_msd_data(path=path if msd_path is None else msd_path, task_name="hepaticvessel", download=download)
100
101    return label_dir, msd_dir
102
103
104def get_hva_ct_paths(
105    path: Union[os.PathLike, str],
106    annotation: Literal["vessels", "liver"] = "vessels",
107    msd_path: Optional[Union[os.PathLike, str]] = None,
108    download: bool = False,
109) -> Tuple[List[str], List[str]]:
110    """Get paths to the HVA-CT data.
111
112    Args:
113        path: Filepath to a folder where the data is downloaded for further processing.
114        annotation: The choice of annotations. Either 'vessels' or 'liver'.
115        msd_path: Filepath to an existing download of the Medical Segmentation Decathlon.
116        download: Whether to download the data if it is not present.
117
118    Returns:
119        List of filepaths for the image data.
120        List of filepaths for the label data.
121    """
122    label_dir, msd_dir = get_hva_ct_data(path, annotation, msd_path, download)
123
124    _, prefix = ANNOTATIONS[annotation]
125    image_paths = _get_image_paths(msd_dir)
126
127    raw_paths, label_paths = [], []
128    for label_path in natsorted(glob(os.path.join(label_dir, "*.nii.gz"))):
129        case_id = os.path.basename(label_path)[len(prefix):-len(".nii.gz")]
130        image_path = image_paths.get(case_id)
131        if image_path is not None:
132            raw_paths.append(image_path)
133            label_paths.append(label_path)
134
135    assert len(raw_paths) == len(label_paths) and len(raw_paths) > 0
136
137    return raw_paths, label_paths
138
139
140def get_hva_ct_dataset(
141    path: Union[os.PathLike, str],
142    patch_shape: Tuple[int, ...],
143    annotation: Literal["vessels", "liver"] = "vessels",
144    msd_path: Optional[Union[os.PathLike, str]] = None,
145    resize_inputs: bool = False,
146    download: bool = False,
147    **kwargs
148) -> Dataset:
149    """Get the HVA-CT dataset for intrahepatic vein and liver segmentation.
150
151    Args:
152        path: Filepath to a folder where the data is downloaded for further processing.
153        patch_shape: The patch shape to use for training.
154        annotation: The choice of annotations. Either 'vessels' or 'liver'.
155        msd_path: Filepath to an existing download of the Medical Segmentation Decathlon.
156        resize_inputs: Whether to resize inputs to the desired patch shape.
157        download: Whether to download the data if it is not present.
158        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
159
160    Returns:
161        The segmentation dataset.
162    """
163    raw_paths, label_paths = get_hva_ct_paths(path, annotation, msd_path, download)
164
165    if resize_inputs:
166        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
167        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
168            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
169        )
170
171    return torch_em.default_segmentation_dataset(
172        raw_paths=raw_paths,
173        raw_key="data",
174        label_paths=label_paths,
175        label_key="data",
176        patch_shape=patch_shape,
177        is_seg_dataset=True,
178        **kwargs
179    )
180
181
182def get_hva_ct_loader(
183    path: Union[os.PathLike, str],
184    batch_size: int,
185    patch_shape: Tuple[int, ...],
186    annotation: Literal["vessels", "liver"] = "vessels",
187    msd_path: Optional[Union[os.PathLike, str]] = None,
188    resize_inputs: bool = False,
189    download: bool = False,
190    **kwargs
191) -> DataLoader:
192    """Get the HVA-CT dataloader for intrahepatic vein and liver segmentation.
193
194    Args:
195        path: Filepath to a folder where the data is downloaded for further processing.
196        batch_size: The batch size for training.
197        patch_shape: The patch shape to use for training.
198        annotation: The choice of annotations. Either 'vessels' or 'liver'.
199        msd_path: Filepath to an existing download of the Medical Segmentation Decathlon.
200        resize_inputs: Whether to resize inputs to the desired patch shape.
201        download: Whether to download the data if it is not present.
202        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
203
204    Returns:
205        The DataLoader.
206    """
207    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
208    dataset = get_hva_ct_dataset(path, patch_shape, annotation, msd_path, resize_inputs, download, **ds_kwargs)
209    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)