torch_em.data.datasets.medical.atlas_liver

  1"""The ATLAS dataset contains annotations for liver and liver tumor segmentation in
  2contrast-enhanced T1-weighted MRI of patients with hepatocellular carcinoma.
  3
  4The training set that is released for the ATLAS challenge consists of 60 CE-MRI volumes with manual
  5delineations of the liver and of the liver tumors. The label ids are: 0 = background, 1 = liver, 2 = tumor.
  6They are taken from the 'dataset.json' of the official release, which also states that the images are T1w.
  7NOTE: The two labels are disjoint, i.e. the liver label covers the parenchyma without the tumors, so a mask of
  8the whole liver is obtained by combining both ids. This was verified on all 60 volumes of the training set.
  9The test set of 30 volumes is not part of the release.
 10
 11NOTE: The dataset is located at https://atlas-challenge.u-bourgogne.fr/dataset and is only available to
 12registered users, so it cannot be downloaded automatically. To download the dataset, please follow these steps:
 13- Visit https://atlas-challenge.u-bourgogne.fr/dataset and create an account via 'Log in / Sign up'.
 14- Log in and download the training set from the 'Get the ATLAS dataset' section.
 15- Place the downloaded zip file (e.g. 'atlas-train-dataset-1.0.1.zip') in the folder passed as 'path'.
 16  This module unzips it and finds the 'imagesTr' and 'labelsTr' folders in the extracted data.
 17
 18The data is licensed under CC BY-NC-SA 4.0.
 19
 20This dataset is from the publication https://doi.org/10.3390/data8050079.
 21Please cite it if you use this dataset in your research.
 22"""
 23
 24import os
 25from glob import glob
 26from natsort import natsorted
 27from typing import Union, Tuple, List
 28
 29from torch.utils.data import Dataset, DataLoader
 30
 31import torch_em
 32
 33from .. import util
 34
 35
 36LABEL_IDS = {"background": 0, "liver": 1, "tumor": 2}
 37
 38
 39def _find_image_dir(path):
 40    """Find the 'imagesTr' folder of the manually downloaded data, which is nested in the release folder."""
 41    image_dirs = glob(os.path.join(path, "**", "imagesTr"), recursive=True)
 42    return image_dirs[0] if image_dirs else None
 43
 44
 45def get_atlas_liver_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 46    """Obtain the ATLAS liver dataset.
 47
 48    Args:
 49        path: Filepath to a folder where the manually downloaded data is stored.
 50        download: Whether to download the data if it is not present. The data cannot be downloaded
 51            automatically, so this raises if the data has not been downloaded manually.
 52
 53    Returns:
 54        Filepath to the folder with the 'imagesTr' and 'labelsTr' folders.
 55    """
 56    image_dir = _find_image_dir(path)
 57    if image_dir is not None:
 58        return os.path.dirname(image_dir)
 59
 60    zip_paths = glob(os.path.join(path, "*atlas*.zip")) + glob(os.path.join(path, "*ATLAS*.zip"))
 61    if not zip_paths:
 62        msg = "'torch_em' cannot download this dataset, because the ATLAS data is only available to users "
 63        msg += "registered at 'https://atlas-challenge.u-bourgogne.fr/dataset'. Please create an account there, "
 64        msg += f"download the training set and place the zip file (e.g. 'atlas-train-dataset-1.0.1.zip') in '{path}'."
 65        raise NotImplementedError(msg)
 66
 67    util.unzip(zip_path=zip_paths[0], dst=path, remove=False)
 68
 69    image_dir = _find_image_dir(path)
 70    if image_dir is None:
 71        raise FileNotFoundError(f"Could not find an 'imagesTr' folder in the data extracted to '{path}'.")
 72
 73    return os.path.dirname(image_dir)
 74
 75
 76def get_atlas_liver_paths(path: Union[os.PathLike, str], download: bool = False) -> Tuple[List[str], List[str]]:
 77    """Get paths to the ATLAS liver data.
 78
 79    Args:
 80        path: Filepath to a folder where the manually downloaded data is stored.
 81        download: Whether to download the data if it is not present.
 82
 83    Returns:
 84        List of filepaths for the image data.
 85        List of filepaths for the label data.
 86    """
 87    data_dir = get_atlas_liver_data(path, download)
 88
 89    raw_paths = natsorted(glob(os.path.join(data_dir, "imagesTr", "*.nii.gz")))
 90    label_paths = natsorted(glob(os.path.join(data_dir, "labelsTr", "*.nii.gz")))
 91    assert len(raw_paths) == len(label_paths) and len(raw_paths) > 0
 92
 93    return raw_paths, label_paths
 94
 95
 96def get_atlas_liver_dataset(
 97    path: Union[os.PathLike, str],
 98    patch_shape: Tuple[int, ...],
 99    resize_inputs: bool = False,
100    download: bool = False,
101    **kwargs
102) -> Dataset:
103    """Get the ATLAS liver dataset for liver and liver tumor segmentation.
104
105    Args:
106        path: Filepath to a folder where the manually downloaded data is stored.
107        patch_shape: The patch shape to use for training.
108        resize_inputs: Whether to resize inputs to the desired patch shape.
109        download: Whether to download the data if it is not present.
110        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
111
112    Returns:
113        The segmentation dataset.
114    """
115    raw_paths, label_paths = get_atlas_liver_paths(path, download)
116
117    if resize_inputs:
118        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
119        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
120            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
121        )
122
123    return torch_em.default_segmentation_dataset(
124        raw_paths=raw_paths,
125        raw_key="data",
126        label_paths=label_paths,
127        label_key="data",
128        patch_shape=patch_shape,
129        is_seg_dataset=True,
130        **kwargs
131    )
132
133
134def get_atlas_liver_loader(
135    path: Union[os.PathLike, str],
136    batch_size: int,
137    patch_shape: Tuple[int, ...],
138    resize_inputs: bool = False,
139    download: bool = False,
140    **kwargs
141) -> DataLoader:
142    """Get the ATLAS liver dataloader for liver and liver tumor segmentation.
143
144    Args:
145        path: Filepath to a folder where the manually downloaded data is stored.
146        batch_size: The batch size for training.
147        patch_shape: The patch shape to use for training.
148        resize_inputs: Whether to resize inputs to the desired patch shape.
149        download: Whether to download the data if it is not present.
150        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
151
152    Returns:
153        The DataLoader.
154    """
155    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
156    dataset = get_atlas_liver_dataset(path, patch_shape, resize_inputs, download, **ds_kwargs)
157    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)