torch_em.data.datasets.medical.duke_liver

  1"""The Duke Liver dataset contains annotations for liver segmentation in MRI scans.
  2
  3NOTE: This dataset is located at https://doi.org/10.5281/zenodo.7774566.
  4Please see 'get_duke_liver_data' for instructions on downloading the dataset.
  5
  6The dataset is from the publication https://doi.org/10.1148/ryai.220275.
  7Please cite it if you use this dataset for your research.
  8"""
  9
 10import os
 11from glob import glob
 12from tqdm import tqdm
 13from natsort import natsorted
 14from typing import Union, Tuple, Literal, List
 15
 16import numpy as np
 17
 18from torch.utils.data import Dataset, DataLoader
 19
 20import torch_em
 21
 22from .. import util
 23
 24
 25def get_duke_liver_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 26    """Get the Duke Liver dataset.
 27
 28    The dataset is located at https://doi.org/10.5281/zenodo.7774566.
 29
 30    Follow the instructions below to get access to the dataset.
 31    - Visit the zenodo site attached above.
 32    - Send a request message alongwith some details to get access to the dataset.
 33    - The authors would accept the request, then you can access the dataset.
 34    - Next, download the `Segmentation.zip` file and provide the path where the zip file is stored.
 35
 36    Args:
 37        path: Filepath to a folder where the data needs to be downloaded for further processing.
 38        download: Whether to download the data if it is not present.
 39
 40    Returns:
 41        Filepath where the data is preprocessed.
 42    """
 43    data_dir = os.path.join(path, "data", "Segmentation")
 44    if os.path.exists(data_dir):
 45        return data_dir
 46
 47    if download:
 48        raise NotImplementedError(
 49            "Automatic download for Duke Liver dataset is not possible. See `get_duke_liver_data` for details."
 50        )
 51
 52    zip_path = os.path.join(path, "Segmentation.zip")
 53    util.unzip(zip_path=zip_path, dst=os.path.join(path, "data"), remove=False)
 54
 55    return data_dir
 56
 57
 58def _preprocess_data(path, data_dir):
 59    preprocess_dir = os.path.join(path, "data", "preprocessed")
 60
 61    if os.path.exists(preprocess_dir):
 62        _image_paths = natsorted(glob(os.path.join(preprocess_dir, "images", "*.nii.gz")))
 63        _gt_paths = natsorted(glob(os.path.join(preprocess_dir, "masks", "*.nii.gz")))
 64        return _image_paths, _gt_paths
 65
 66    os.makedirs(os.path.join(preprocess_dir, "images"), exist_ok=True)
 67    os.makedirs(os.path.join(preprocess_dir, "masks"), exist_ok=True)
 68
 69    image_paths, gt_paths = [], []
 70    for patient_dir in tqdm(glob(os.path.join(data_dir, "00*"))):
 71        patient_id = os.path.split(patient_dir)[-1]
 72
 73        for sub_id_dir in glob(os.path.join(patient_dir, "*")):
 74            sub_id = os.path.split(sub_id_dir)[-1]
 75
 76            image_path = os.path.join(preprocess_dir, "images", f"{patient_id}_{sub_id}.nii.gz")
 77            gt_path = os.path.join(preprocess_dir, "masks", f"{patient_id}_{sub_id}.nii.gz")
 78
 79            image_paths.append(image_path)
 80            gt_paths.append(gt_path)
 81
 82            if os.path.exists(image_path) and os.path.exists(gt_path):
 83                continue
 84
 85            image_slice_paths = natsorted(glob(os.path.join(sub_id_dir, "images", "*.dicom")))
 86            gt_slice_paths = natsorted(glob(os.path.join(sub_id_dir, "masks", "*.dicom")))
 87
 88            import pydicom as dicom
 89            import nibabel as nib
 90
 91            images, gts = [], []
 92            for image_slice_path, gt_slice_path in zip(image_slice_paths, gt_slice_paths):
 93                image_slice = dicom.dcmread(image_slice_path).pixel_array
 94                gt_slice = dicom.dcmread(gt_slice_path).pixel_array
 95
 96                images.append(image_slice)
 97                gts.append(gt_slice)
 98
 99            image = np.stack(images).transpose(1, 2, 0)
100            gt = np.stack(gts).transpose(1, 2, 0)
101
102            assert image.shape == gt.shape
103
104            image = nib.Nifti2Image(image, np.eye(4))
105            gt = nib.Nifti2Image(gt, np.eye(4))
106
107            nib.save(image, image_path)
108            nib.save(gt, gt_path)
109
110    return natsorted(image_paths), natsorted(gt_paths)
111
112
113def get_duke_liver_paths(
114    path: Union[os.PathLike, str], split: Literal['train', 'val', 'test'], download: bool = False
115) -> Tuple[List[str], List[str]]:
116    """Get paths to the Duke Liver dataset.
117
118    Args:
119        path: Filepath to a folder where the data needs to be downloaded for further processing.
120        split: The choice of data split.
121        download: Whether to download the data if it is not present.
122
123    Returns:
124        List of filepaths for the image data.
125        List of filepaths for the label data.
126    """
127    data_dir = get_duke_liver_data(path=path, download=download)
128
129    image_paths, gt_paths = _preprocess_data(path=path, data_dir=data_dir)
130
131    if split == "train":
132        image_paths, gt_paths = image_paths[:250], gt_paths[:250]
133    elif split == "val":
134        image_paths, gt_paths = image_paths[250:260], gt_paths[250:260]
135    elif split == "test":
136        image_paths, gt_paths = image_paths[260:], gt_paths[260:]
137    else:
138        raise ValueError(f"'{split}' is not a valid split.")
139
140    return image_paths, gt_paths
141
142
143def get_duke_liver_dataset(
144    path: Union[os.PathLike, str],
145    patch_shape: Tuple[int, ...],
146    split: Literal['train', 'val', 'test'],
147    resize_inputs: bool = False,
148    download: bool = False,
149    **kwargs
150) -> Dataset:
151    """Get the Duke Liver dataset for segmentation of liver in MRI.
152
153    Args:
154        path: Filepath to a folder where the data needs to be downloaded for further processing.
155        patch_shape: The patch shape to use for training.
156        split: The choice of data split.
157        resize_inputs: Whether to resize the inputs to the patch shape.
158        download: Whether to download the data if it is not present.
159        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
160
161    Returns:
162        The segmentation dataset.
163    """
164    image_paths, gt_paths = get_duke_liver_paths(path, split, download)
165
166    if resize_inputs:
167        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
168        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
169            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
170        )
171
172    return torch_em.default_segmentation_dataset(
173        raw_paths=image_paths,
174        raw_key="data",
175        label_paths=gt_paths,
176        label_key="data",
177        is_seg_dataset=True,
178        patch_shape=patch_shape,
179        **kwargs
180    )
181
182
183def get_duke_liver_loader(
184    path: Union[os.PathLike, str],
185    batch_size: int,
186    patch_shape: Tuple[int, ...],
187    split: Literal['train', 'val', 'test'],
188    resize_inputs: bool = False,
189    download: bool = False,
190    **kwargs
191) -> DataLoader:
192    """Get the Duke Liver dataloader for segmentation of liver in MRI.
193
194    Args:
195        path: Filepath to a folder where the data needs to be downloaded for further processing.
196        batch_size: The batch size for training.
197        patch_shape: The patch shape to use for training.
198        split: The choice of data split.
199        resize_inputs: Whether to resize the inputs to the patch shape.
200        download: Whether to download the data if it is not present.
201        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
202
203    Returns:
204        The DataLoader.
205    """
206    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
207    dataset = get_duke_liver_dataset(path, patch_shape, split, resize_inputs, download, **ds_kwargs)
208    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)