torch_em.data.datasets.medical.sliver07

  1"""The SLIVER07 dataset contains annotations for liver segmentation in CT scans.
  2
  3The dataset consists of the 20 training CT scans of the MICCAI 2007 liver segmentation challenge with a
  4binary liver mask. The scans are distributed as MetaImage files, which this module converts into nifti
  5files once. The 10 test scans of the challenge are not distributed with labels.
  6
  7The dataset is located at https://doi.org/10.5281/zenodo.2597908. Its license restricts use to
  8non-commercial liver segmentation research and requires the publication below to be cited; see the
  9'license.txt' of the release for the complete terms.
 10This dataset is from the publication https://doi.org/10.1109/TMI.2009.2013851.
 11Please cite it if you use this dataset in your research.
 12"""
 13
 14import os
 15from glob import glob
 16from tqdm import tqdm
 17from natsort import natsorted
 18from typing import Union, Tuple, List
 19
 20import numpy as np
 21
 22from torch.utils.data import Dataset, DataLoader
 23
 24import torch_em
 25
 26from .mediastinal_ct import read_mhd
 27from .. import util
 28
 29
 30URLS = {
 31    "training-scans": "https://zenodo.org/records/2597908/files/training-scans.zip?download=1",
 32    "training-labels": "https://zenodo.org/records/2597908/files/training-labels.zip?download=1",
 33}
 34
 35CHECKSUMS = {
 36    "training-scans": "caf4e43650050543ec59a07c2cb983faa023696dbe3f958019bfa025f13dc6c3",
 37    "training-labels": "72c41c299ce1392cf810424268186518157c14c361e7dc2220657b1cf7173d78",
 38}
 39
 40
 41def _convert_sliver07_to_nifti(data_dir, preprocessed_dir):
 42    import nibabel as nib
 43
 44    os.makedirs(preprocessed_dir, exist_ok=True)
 45    image_paths = natsorted(glob(os.path.join(data_dir, "**", "liver-orig*.mhd"), recursive=True))
 46    for image_path in tqdm(image_paths, desc="Converting SLIVER07 to nifti"):
 47        scan_id = os.path.basename(image_path)[len("liver-orig"):-len(".mhd")]
 48        label_paths = glob(os.path.join(data_dir, "**", f"liver-seg{scan_id}.mhd"), recursive=True)
 49        if not label_paths:
 50            continue
 51
 52        out_image_path = os.path.join(preprocessed_dir, f"liver-orig{scan_id}.nii.gz")
 53        out_label_path = os.path.join(preprocessed_dir, f"liver-seg{scan_id}.nii.gz")
 54        if os.path.exists(out_image_path) and os.path.exists(out_label_path):
 55            continue
 56
 57        volume, spacing = read_mhd(image_path)
 58        labels, _ = read_mhd(label_paths[0])
 59        assert labels.shape == volume.shape, f"The mask of scan '{scan_id}' does not match its scan."
 60
 61        affine = np.diag(list(spacing) + [1.0])
 62        nib.save(nib.Nifti1Image(volume, affine), out_image_path)
 63        nib.save(nib.Nifti1Image(labels.astype("uint8"), affine), out_label_path)
 64
 65
 66def get_sliver07_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 67    """Download the SLIVER07 dataset.
 68
 69    Args:
 70        path: Filepath to a folder where the data is downloaded for further processing.
 71        download: Whether to download the data if it is not present.
 72
 73    Returns:
 74        Filepath where the preprocessed data is stored.
 75    """
 76    # NOTE: The conversion below skips scans that were converted already, so an interrupted run resumes.
 77    preprocessed_dir = os.path.join(path, "preprocessed")
 78
 79    if not glob(os.path.join(path, "**", "liver-orig*.mhd"), recursive=True):
 80        os.makedirs(path, exist_ok=True)
 81        for name, url in URLS.items():
 82            zip_path = os.path.join(path, f"{name}.zip")
 83            util.download_source(path=zip_path, url=url, download=download, checksum=CHECKSUMS[name])
 84            util.unzip(zip_path=zip_path, dst=path, remove=False)
 85
 86    _convert_sliver07_to_nifti(path, preprocessed_dir)
 87    return preprocessed_dir
 88
 89
 90def get_sliver07_paths(
 91    path: Union[os.PathLike, str], download: bool = False
 92) -> Tuple[List[str], List[str]]:
 93    """Get paths to the SLIVER07 data.
 94
 95    Args:
 96        path: Filepath to a folder where the data is downloaded for further processing.
 97        download: Whether to download the data if it is not present.
 98
 99    Returns:
100        List of filepaths for the image data.
101        List of filepaths for the label data.
102    """
103    preprocessed_dir = get_sliver07_data(path, download)
104
105    raw_paths = natsorted(glob(os.path.join(preprocessed_dir, "liver-orig*.nii.gz")))
106    label_paths = [p.replace("liver-orig", "liver-seg") for p in raw_paths]
107
108    assert len(raw_paths) == len(label_paths) and len(raw_paths) > 0
109    assert all(os.path.exists(p) for p in label_paths)
110
111    return raw_paths, label_paths
112
113
114def get_sliver07_dataset(
115    path: Union[os.PathLike, str],
116    patch_shape: Tuple[int, ...],
117    resize_inputs: bool = False,
118    download: bool = False,
119    **kwargs
120) -> Dataset:
121    """Get the SLIVER07 dataset for liver segmentation.
122
123    Args:
124        path: Filepath to a folder where the data is downloaded for further processing.
125        patch_shape: The patch shape to use for training.
126        resize_inputs: Whether to resize inputs to the desired patch shape.
127        download: Whether to download the data if it is not present.
128        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
129
130    Returns:
131        The segmentation dataset.
132    """
133    raw_paths, label_paths = get_sliver07_paths(path, download)
134
135    if resize_inputs:
136        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
137        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
138            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
139        )
140
141    return torch_em.default_segmentation_dataset(
142        raw_paths=raw_paths,
143        raw_key="data",
144        label_paths=label_paths,
145        label_key="data",
146        patch_shape=patch_shape,
147        is_seg_dataset=True,
148        **kwargs
149    )
150
151
152def get_sliver07_loader(
153    path: Union[os.PathLike, str],
154    batch_size: int,
155    patch_shape: Tuple[int, ...],
156    resize_inputs: bool = False,
157    download: bool = False,
158    **kwargs
159) -> DataLoader:
160    """Get the SLIVER07 dataloader for liver segmentation.
161
162    Args:
163        path: Filepath to a folder where the data is downloaded for further processing.
164        batch_size: The batch size for training.
165        patch_shape: The patch shape to use for training.
166        resize_inputs: Whether to resize inputs to the desired patch shape.
167        download: Whether to download the data if it is not present.
168        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
169
170    Returns:
171        The DataLoader.
172    """
173    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
174    dataset = get_sliver07_dataset(path, patch_shape, resize_inputs, download, **ds_kwargs)
175    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)