torch_em.data.datasets.medical.couinaud

  1"""The Couinaud dataset contains annotations for the eight Couinaud liver segments in CT scans.
  2
  3The annotations were made on the CT scans of the Medical Segmentation Decathlon task 8 (hepatic vessel),
  4so the images are downloaded from there. Two annotation sets are provided: 'couinaud' with the eight
  5Couinaud segments for 193 scans and 'liver' with a binary liver mask for 443 scans. The label ids of the
  6'couinaud' annotations are 1 to 8 for the Couinaud segments I to VIII. See also `CLASS_IDS`.
  7
  8The dataset is located at https://github.com/GLCUnet/dataset and is distributed under the MIT license.
  9This dataset is from the publication https://doi.org/10.1007/978-3-030-32692-0_32.
 10Please cite it if you use this dataset in your research.
 11"""
 12
 13import os
 14from glob import glob
 15from natsort import natsorted
 16from typing import Union, Tuple, Literal, List
 17
 18from torch.utils.data import Dataset, DataLoader
 19
 20import torch_em
 21
 22from .msd import get_msd_data
 23from .. import util
 24
 25
 26URLS = {
 27    "couinaud": "https://raw.githubusercontent.com/GLCUnet/dataset/master/couinaud_annotation.zip",
 28    "liver": "https://raw.githubusercontent.com/GLCUnet/dataset/master/liver_annotation.zip",
 29}
 30
 31CHECKSUMS = {
 32    "couinaud": "fb2fc7809a7982267adc2785dddd5af3770070fea230e2f01142fe799a29f0cf",
 33    "liver": "e14148cec317829a1d4068e82f60cfecea6df570decd7c06b474f172aff382a1",
 34}
 35
 36CLASS_NAMES = [
 37    "segment_i", "segment_ii", "segment_iii", "segment_iv",
 38    "segment_v", "segment_vi", "segment_vii", "segment_viii",
 39]
 40"""The Couinaud liver segments. The label id of a segment is its 1-based index."""
 41
 42CLASS_IDS = {name: i + 1 for i, name in enumerate(CLASS_NAMES)}
 43"""Mapping from the Couinaud segment name to its label id."""
 44
 45
 46def _get_image_paths(msd_dir):
 47    image_paths = {}
 48    for split in ["imagesTr", "imagesTs"]:
 49        for path in glob(os.path.join(msd_dir, "Task08_HepaticVessel", split, "*.nii.gz")):
 50            fname = os.path.basename(path)
 51            # The MSD archives carry macOS resource fork files next to the actual volumes.
 52            if fname.startswith("._"):
 53                continue
 54            image_paths[fname] = path
 55    return image_paths
 56
 57
 58def get_couinaud_data(
 59    path: Union[os.PathLike, str], annotation: Literal["couinaud", "liver"] = "couinaud", download: bool = False
 60) -> Tuple[str, str]:
 61    """Download the Couinaud dataset.
 62
 63    Args:
 64        path: Filepath to a folder where the data is downloaded for further processing.
 65        annotation: The choice of annotations. Either 'couinaud' for the eight Couinaud segments
 66            or 'liver' for a binary liver mask.
 67        download: Whether to download the data if it is not present.
 68
 69    Returns:
 70        Filepath where the annotations are downloaded.
 71        Filepath where the images are downloaded.
 72    """
 73    if annotation not in URLS:
 74        raise ValueError(f"'{annotation}' is not a valid annotation. Choose from {list(URLS.keys())}.")
 75
 76    label_dir = os.path.join(path, f"{annotation}_annotation")
 77    if not os.path.exists(label_dir):
 78        os.makedirs(path, exist_ok=True)
 79        zip_path = os.path.join(path, f"{annotation}_annotation.zip")
 80        util.download_source(
 81            path=zip_path, url=URLS[annotation], download=download, checksum=CHECKSUMS[annotation]
 82        )
 83        util.unzip(zip_path=zip_path, dst=path, remove=False)
 84
 85    # The images are the hepatic vessel scans of the Medical Segmentation Decathlon.
 86    msd_dir = get_msd_data(path=path, task_name="hepaticvessel", download=download)
 87
 88    return label_dir, msd_dir
 89
 90
 91def get_couinaud_paths(
 92    path: Union[os.PathLike, str],
 93    annotation: Literal["couinaud", "liver"] = "couinaud",
 94    download: bool = False,
 95) -> Tuple[List[str], List[str]]:
 96    """Get paths to the Couinaud data.
 97
 98    Args:
 99        path: Filepath to a folder where the data is downloaded for further processing.
100        annotation: The choice of annotations. Either 'couinaud' for the eight Couinaud segments
101            or 'liver' for a binary liver mask.
102        download: Whether to download the data if it is not present.
103
104    Returns:
105        List of filepaths for the image data.
106        List of filepaths for the label data.
107    """
108    label_dir, msd_dir = get_couinaud_data(path, annotation, download)
109
110    image_paths = _get_image_paths(msd_dir)
111    label_paths = natsorted(glob(os.path.join(label_dir, "*.nii.gz")))
112
113    raw_paths, valid_label_paths = [], []
114    for label_path in label_paths:
115        image_path = image_paths.get(os.path.basename(label_path))
116        if image_path is not None:
117            raw_paths.append(image_path)
118            valid_label_paths.append(label_path)
119
120    assert len(raw_paths) == len(valid_label_paths) and len(raw_paths) > 0
121
122    return raw_paths, valid_label_paths
123
124
125def get_couinaud_dataset(
126    path: Union[os.PathLike, str],
127    patch_shape: Tuple[int, ...],
128    annotation: Literal["couinaud", "liver"] = "couinaud",
129    resize_inputs: bool = False,
130    download: bool = False,
131    **kwargs
132) -> Dataset:
133    """Get the Couinaud dataset for liver segment segmentation.
134
135    Args:
136        path: Filepath to a folder where the data is downloaded for further processing.
137        patch_shape: The patch shape to use for training.
138        annotation: The choice of annotations. Either 'couinaud' for the eight Couinaud segments
139            or 'liver' for a binary liver mask.
140        resize_inputs: Whether to resize inputs to the desired patch shape.
141        download: Whether to download the data if it is not present.
142        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
143
144    Returns:
145        The segmentation dataset.
146    """
147    raw_paths, label_paths = get_couinaud_paths(path, annotation, download)
148
149    if resize_inputs:
150        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
151        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
152            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
153        )
154
155    return torch_em.default_segmentation_dataset(
156        raw_paths=raw_paths,
157        raw_key="data",
158        label_paths=label_paths,
159        label_key="data",
160        patch_shape=patch_shape,
161        is_seg_dataset=True,
162        **kwargs
163    )
164
165
166def get_couinaud_loader(
167    path: Union[os.PathLike, str],
168    batch_size: int,
169    patch_shape: Tuple[int, ...],
170    annotation: Literal["couinaud", "liver"] = "couinaud",
171    resize_inputs: bool = False,
172    download: bool = False,
173    **kwargs
174) -> DataLoader:
175    """Get the Couinaud dataloader for liver segment segmentation.
176
177    Args:
178        path: Filepath to a folder where the data is downloaded for further processing.
179        batch_size: The batch size for training.
180        patch_shape: The patch shape to use for training.
181        annotation: The choice of annotations. Either 'couinaud' for the eight Couinaud segments
182            or 'liver' for a binary liver mask.
183        resize_inputs: Whether to resize inputs to the desired patch shape.
184        download: Whether to download the data if it is not present.
185        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
186
187    Returns:
188        The DataLoader.
189    """
190    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
191    dataset = get_couinaud_dataset(path, patch_shape, annotation, resize_inputs, download, **ds_kwargs)
192    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)