torch_em.data.datasets.medical.waw_tace

  1"""The WAW-TACE dataset contains annotations for liver tumors in multiphase CT scans.
  2
  3The dataset consists of multiphase abdominal CT scans of 233 treatment-naive patients with hepatocellular
  4carcinoma that were treated with transarterial chemoembolization. The tumors of one phase per patient are
  5delineated by hand, giving 378 masks. The masks are stored as instance labels, so that the tumors of a
  6scan get the ids 1 to n.
  7
  8NOTE: The release also holds masks of several internal organs, but those were generated with
  9TotalSegmentator rather than drawn by hand, so they are predictions and not annotations and are not
 10provided here. `medical.totalsegmentator` provides the data that model was trained on.
 11
 12NOTE: This requires the pynrrd python package to read the tumor masks.
 13
 14The dataset is located at https://doi.org/10.5281/zenodo.12741586 and is distributed under the
 15CC BY 4.0 license.
 16This dataset is from the publication https://doi.org/10.1148/ryai.240296.
 17Please cite it if you use this dataset in your research.
 18"""
 19
 20import os
 21from glob import glob
 22from tqdm import tqdm
 23from natsort import natsorted
 24from collections import defaultdict
 25from typing import Union, Tuple, List
 26
 27import numpy as np
 28
 29from torch.utils.data import Dataset, DataLoader
 30
 31import torch_em
 32
 33from .. import util
 34
 35
 36URL = "https://zenodo.org/records/12741586/files/{filename}?download=1"
 37
 38SCAN_ARCHIVES = [f"ct_scans_{index}_4_wawtace_09_05_24.zip" for index in range(1, 5)]
 39
 40TUMOR_ARCHIVE = "tumor_masks_wawtace_v1_08_05_2024.zip"
 41
 42CHECKSUMS = {TUMOR_ARCHIVE: "39268ef2899cb5ccd422950d7f352a173c0f54071a5e2775157501e4756dff1e"}
 43
 44
 45def _mask_geometry(mask_path):
 46    """The shape, origin and spacing of a mask, with the signs of its LPS origin dropped."""
 47    import nrrd
 48
 49    header = nrrd.read_header(mask_path)
 50    origin = np.abs(np.array([float(value) for value in header["space origin"]]))
 51    spacing = np.abs(np.diag(np.array(header["space directions"], dtype="float64")))
 52    return tuple(int(size) for size in header["sizes"]), tuple(np.round(origin, 2)), tuple(np.round(spacing, 3))
 53
 54
 55def _scan_geometry(image_path):
 56    """The same geometry for a scan. Its nifti axes are LAS, so the signs of the origin differ from the mask."""
 57    import nibabel as nib
 58
 59    image = nib.load(image_path)
 60    affine = image.affine
 61    return (
 62        tuple(image.shape),
 63        tuple(np.round(np.abs(affine[:3, 3]), 2)),
 64        tuple(np.round(np.abs(np.diag(affine[:3, :3])), 3)),
 65    )
 66
 67
 68def _find_scan(data_dir, patient_id, phase, mask_path):
 69    """Find the scan a mask was drawn on.
 70
 71    The phase in the name of a mask is not always the phase of the scan it belongs to, so the geometry
 72    decides. The named phase is preferred, because the phases of a study often share their geometry and
 73    the match would otherwise be ambiguous.
 74    """
 75    geometry = _mask_geometry(mask_path)
 76    named = os.path.join(data_dir, patient_id, f"{patient_id}_{phase}_scan.nii.gz")
 77    if os.path.exists(named) and _scan_geometry(named) == geometry:
 78        return named
 79
 80    matches = [
 81        path for path in natsorted(glob(os.path.join(data_dir, patient_id, f"{patient_id}_*_scan.nii.gz")))
 82        if _scan_geometry(path) == geometry
 83    ]
 84    return matches[0] if len(matches) == 1 else None
 85
 86
 87def _preprocess_waw_tace(data_dir, tumor_dir, preprocessed_dir):
 88    import h5py
 89    import nrrd
 90    import nibabel as nib
 91
 92    # The tumors of a scan are stored one per file, as '<patient>_<phase>_<tumor>_tumor_seg.nrrd'.
 93    tumors = defaultdict(list)
 94    for mask_path in natsorted(glob(os.path.join(tumor_dir, "*", "*_tumor_seg.nrrd"))):
 95        patient_id, phase = os.path.basename(mask_path).split("_")[:2]
 96        tumors[(patient_id, phase)].append(mask_path)
 97
 98    os.makedirs(preprocessed_dir, exist_ok=True)
 99    for (patient_id, phase), mask_paths in tqdm(sorted(tumors.items()), desc="Preprocess WAW-TACE"):
100        out_path = os.path.join(preprocessed_dir, f"{patient_id}_{phase}.h5")
101        if os.path.exists(out_path):
102            continue
103
104        image_path = _find_scan(data_dir, patient_id, phase, mask_paths[0])
105        if image_path is None:
106            continue
107
108        image = nib.load(image_path)
109        # The scans and their masks are stored with axis order (x, y, z) and are transposed to (z, y, x).
110        volume = np.asarray(image.dataobj).transpose(2, 1, 0)
111        labels = np.zeros(volume.shape, dtype="uint8")
112        for instance_id, mask_path in enumerate(mask_paths, start=1):
113            mask, _ = nrrd.read(mask_path)
114            mask = mask.transpose(2, 1, 0)
115            if mask.shape != volume.shape:
116                continue
117            labels[mask > 0] = instance_id
118
119        if labels.max() == 0:
120            continue
121
122        with h5py.File(out_path, "w") as f:
123            f.create_dataset("raw", data=volume, compression="gzip")
124            f.create_dataset("labels", data=labels, compression="gzip")
125
126
127def get_waw_tace_data(path: Union[os.PathLike, str], download: bool = False) -> str:
128    """Download the WAW-TACE dataset.
129
130    Args:
131        path: Filepath to a folder where the data is downloaded for further processing.
132        download: Whether to download the data if it is not present.
133
134    Returns:
135        Filepath where the preprocessed data is stored.
136    """
137    # NOTE: The preprocessing below skips volumes that were converted already, so an interrupted run resumes.
138    preprocessed_dir = os.path.join(path, "preprocessed")
139
140    os.makedirs(path, exist_ok=True)
141    data_dir = os.path.join(path, "scans")
142    for filename in SCAN_ARCHIVES:
143        zip_path = os.path.join(path, filename)
144        util.download_source(path=zip_path, url=URL.format(filename=filename), download=download, checksum=None)
145        util.unzip(zip_path=zip_path, dst=data_dir, remove=False)
146
147    tumor_dir = os.path.join(path, TUMOR_ARCHIVE[:-len(".zip")])
148    if not os.path.exists(tumor_dir):
149        zip_path = os.path.join(path, TUMOR_ARCHIVE)
150        util.download_source(
151            path=zip_path, url=URL.format(filename=TUMOR_ARCHIVE), download=download,
152            checksum=CHECKSUMS[TUMOR_ARCHIVE],
153        )
154        util.unzip(zip_path=zip_path, dst=path, remove=False)
155
156    _preprocess_waw_tace(data_dir, tumor_dir, preprocessed_dir)
157    return preprocessed_dir
158
159
160def get_waw_tace_paths(path: Union[os.PathLike, str], download: bool = False) -> List[str]:
161    """Get paths to the WAW-TACE data.
162
163    Args:
164        path: Filepath to a folder where the data is downloaded for further processing.
165        download: Whether to download the data if it is not present.
166
167    Returns:
168        List of filepaths for the stored data.
169    """
170    preprocessed_dir = get_waw_tace_data(path, download)
171    volume_paths = natsorted(glob(os.path.join(preprocessed_dir, "*.h5")))
172    assert len(volume_paths) > 0, f"Could not find any preprocessed volumes in '{preprocessed_dir}'."
173    return volume_paths
174
175
176def get_waw_tace_dataset(
177    path: Union[os.PathLike, str],
178    patch_shape: Tuple[int, ...],
179    resize_inputs: bool = False,
180    download: bool = False,
181    **kwargs
182) -> Dataset:
183    """Get the WAW-TACE dataset for liver tumor segmentation.
184
185    Args:
186        path: Filepath to a folder where the data is downloaded for further processing.
187        patch_shape: The patch shape to use for training.
188        resize_inputs: Whether to resize inputs to the desired patch shape.
189        download: Whether to download the data if it is not present.
190        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
191
192    Returns:
193        The segmentation dataset.
194    """
195    volume_paths = get_waw_tace_paths(path, download)
196
197    if resize_inputs:
198        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
199        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
200            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
201        )
202
203    return torch_em.default_segmentation_dataset(
204        raw_paths=volume_paths,
205        raw_key="raw",
206        label_paths=volume_paths,
207        label_key="labels",
208        patch_shape=patch_shape,
209        is_seg_dataset=True,
210        **kwargs
211    )
212
213
214def get_waw_tace_loader(
215    path: Union[os.PathLike, str],
216    batch_size: int,
217    patch_shape: Tuple[int, ...],
218    resize_inputs: bool = False,
219    download: bool = False,
220    **kwargs
221) -> DataLoader:
222    """Get the WAW-TACE dataloader for liver tumor segmentation.
223
224    Args:
225        path: Filepath to a folder where the data is downloaded for further processing.
226        batch_size: The batch size for training.
227        patch_shape: The patch shape to use for training.
228        resize_inputs: Whether to resize inputs to the desired patch shape.
229        download: Whether to download the data if it is not present.
230        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
231
232    Returns:
233        The DataLoader.
234    """
235    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
236    dataset = get_waw_tace_dataset(path, patch_shape, resize_inputs, download, **ds_kwargs)
237    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)