torch_em.data.datasets.medical.deeplesion

  1"""The DeepLesion3D dataset contains 3D lesion annotations for CT volumes-of-interest from the DeepLesion dataset.
  2
  3The original DeepLesion dataset (https://nihcc.app.box.com/v/DeepLesion) only provides RECIST diameters and 2D
  4bounding boxes on the key slice of each lesion. This module provides the 'ULS23_DeepLesion3D' subset that was
  5created for the Universal Lesion Segmentation Challenge 2023 (ULS23): 750 lesions from DeepLesion were segmented
  6in 3D by trained (bio-)medical students (each lesion in triplicate, the majority vote is the final label).
  7Each sample is a 256 x 256 x 128 volume-of-interest cropped around one lesion, with a binary label mask (1: lesion).
  8The volumes are padded with a constant value along the z-axis, so that the lesion is centered around slice 64.
  9743 of the 750 volumes have a label; the remaining 7 are skipped. The lesions are grouped into 7 categories:
 10200 abdominal, 100 bone, 50 kidney, 50 liver, 100 lung, 100 mediastinal and 150 other lesions.
 11
 12The images are located at https://doi.org/10.5281/zenodo.10035161 (part 1 of the ULS23 training data, which also
 13contains the bone and pancreas lesions from Radboudumc).
 14The labels are located at https://github.com/DIAGNijmegen/ULS23.
 15NOTE: The images are distributed as a multi-part zip archive, so the '7z' CLI is required to extract them
 16(install it via 'conda install -c conda-forge p7zip').
 17The data is licensed under CC BY-NC-SA 4.0.
 18
 19This dataset is from the publication https://doi.org/10.1016/j.media.2025.103525.
 20Please cite it (and the DeepLesion publication https://doi.org/10.1117/1.JMI.5.3.036501) if you use this dataset
 21in your research.
 22"""
 23
 24import os
 25import json
 26from glob import glob
 27from tqdm import tqdm
 28from shutil import which, rmtree
 29from subprocess import run
 30from natsort import natsorted
 31from typing import Union, Tuple, Literal, List, Optional
 32
 33import numpy as np
 34
 35from torch.utils.data import Dataset, DataLoader
 36
 37import torch_em
 38
 39from .. import util
 40
 41
 42URLS = {
 43    "ULS23_Part1.zip": "https://zenodo.org/records/10035161/files/ULS23_Part1.zip?download=1",
 44    "ULS23_Part1.z01": "https://zenodo.org/records/10035161/files/ULS23_Part1.z01?download=1",
 45    "ULS23_Part1.z02": "https://zenodo.org/records/10035161/files/ULS23_Part1.z02?download=1",
 46    "ULS23_Part1.z03": "https://zenodo.org/records/10035161/files/ULS23_Part1.z03?download=1",
 47    "annotations": "https://github.com/DIAGNijmegen/ULS23/archive/06a2bffc433418f72d04f7ecbb23b28694c81e6b.zip",
 48}
 49
 50CHECKSUMS = {
 51    "ULS23_Part1.zip": "b0beff9cedd09b212087f8fb678e07de8c41126ba11ee2c2ad14d7e210b8fabb",
 52    "ULS23_Part1.z01": "d23eefdfec9c394f9ce7effb0daa3c9dac37d38802c063eed2aa5e2a6b7aa354",
 53    "ULS23_Part1.z02": "d0b464b5ff0b099a85809ed354203bab2830b90208d17df6b30344dc277a519c",
 54    "ULS23_Part1.z03": "12f2ecef0788267a1f96b21e413647e15014cc8a2f40191dfa5d55426c337200",
 55    "annotations": "19ae6b84aae1a94aa8329a0ce4c6586d1e7cf28532b27367d923773adc3ebe28",
 56}
 57
 58CATEGORIES = ["Abdominal", "Bone", "Kidney", "Liver", "Lung", "Mediastinal", "Other"]
 59"""The lesion categories of the DeepLesion3D dataset."""
 60
 61ULS_DIR = os.path.join("ULS23", "novel_data", "ULS23_DeepLesion3D")
 62
 63
 64def _extract_uls_archives(path, download):
 65    zip_path = os.path.join(path, "ULS23_Part1.zip")
 66    for name, url in URLS.items():
 67        if name == "annotations":
 68            continue
 69        util.download_source(path=os.path.join(path, name), url=url, download=download, checksum=CHECKSUMS[name])
 70
 71    if which("7z") is None:
 72        raise RuntimeError(
 73            "The DeepLesion3D images are distributed as a multi-part zip archive, which requires the '7z' CLI. "
 74            "You can install it via 'conda install -c conda-forge p7zip'."
 75        )
 76
 77    inner_zips = [os.path.join(ULS_DIR, "images.zip"), os.path.join(ULS_DIR, "categories.zip")]
 78    run(["7z", "x", f"-o{path}", "-y", zip_path] + inner_zips, check=True)
 79    for inner_zip in inner_zips:
 80        util.unzip(zip_path=os.path.join(path, inner_zip), dst=os.path.join(path, ULS_DIR), remove=True)
 81
 82    annotation_zip = os.path.join(path, "ULS23_annotations.zip")
 83    util.download_source(
 84        path=annotation_zip, url=URLS["annotations"], download=download, checksum=CHECKSUMS["annotations"]
 85    )
 86    util.unzip(zip_path=annotation_zip, dst=os.path.join(path, "ULS23_annotations"), remove=True)
 87
 88
 89def _preprocess_deeplesion3d(path, data_dir):
 90    import nibabel as nib
 91
 92    image_dir = os.path.join(data_dir, "images")
 93    label_dir = os.path.join(data_dir, "labels")
 94    os.makedirs(image_dir, exist_ok=True)
 95    os.makedirs(label_dir, exist_ok=True)
 96
 97    src_image_dir = os.path.join(path, ULS_DIR, "images")
 98    category_dir = os.path.join(path, ULS_DIR, "categories")
 99    annotation_dir = os.path.join(path, "ULS23_annotations", "*", "annotations", ULS_DIR, "labels")
100    label_zips = glob(os.path.join(annotation_dir, "*.zip"))
101    assert len(label_zips) > 0, "Could not find the DeepLesion3D annotations."
102
103    categories = {}
104    for category in CATEGORIES:
105        for fname in os.listdir(os.path.join(category_dir, category)):
106            categories[os.path.splitext(fname)[0]] = category
107
108    for label_zip in tqdm(natsorted(label_zips), desc="Preprocessing DeepLesion3D"):
109        fname = os.path.basename(label_zip)[:-len(".zip")]
110        src_image_path = os.path.join(src_image_dir, fname)
111        assert os.path.exists(src_image_path), src_image_path
112
113        util.unzip(zip_path=label_zip, dst=label_dir, remove=False)
114        label_path = os.path.join(label_dir, fname)
115        image_path = os.path.join(image_dir, fname)
116
117        # The volumes are stored with a trailing singleton dimension, which we remove.
118        for src, dst, dtype in [(src_image_path, image_path, "float32"), (label_path, label_path, "uint8")]:
119            nifti = nib.load(src)
120            data = np.squeeze(np.asarray(nifti.dataobj)).astype(dtype)
121            assert data.ndim == 3, data.shape
122            nib.save(nib.Nifti1Image(data, nifti.affine), dst)
123
124    # The category of a lesion is given by the first four fields of its name (patient, study, series, key slice).
125    lesion_categories = {}
126    for label_path in glob(os.path.join(label_dir, "*.nii.gz")):
127        fname = os.path.basename(label_path)
128        lesion_categories[fname] = categories["_".join(fname.split("_")[:4])]
129    with open(os.path.join(data_dir, "categories.json"), "w") as f:
130        json.dump(lesion_categories, f, indent=2, sort_keys=True)
131
132    # Remove the intermediate data.
133    rmtree(os.path.join(path, "ULS23"))
134    rmtree(os.path.join(path, "ULS23_annotations"))
135
136
137def get_deeplesion_data(path: Union[os.PathLike, str], download: bool = False) -> str:
138    """Download the DeepLesion3D dataset.
139
140    Args:
141        path: Filepath to a folder where the data is downloaded for further processing.
142        download: Whether to download the data if it is not present.
143
144    Returns:
145        Filepath where the data is downloaded.
146    """
147    data_dir = os.path.join(path, "DeepLesion3D")
148    if os.path.exists(os.path.join(data_dir, "categories.json")):
149        return data_dir
150
151    os.makedirs(path, exist_ok=True)
152    _extract_uls_archives(path, download)
153    _preprocess_deeplesion3d(path, data_dir)
154
155    return data_dir
156
157
158def get_deeplesion_paths(
159    path: Union[os.PathLike, str],
160    category: Optional[Literal["Abdominal", "Bone", "Kidney", "Liver", "Lung", "Mediastinal", "Other"]] = None,
161    download: bool = False,
162) -> Tuple[List[str], List[str]]:
163    """Get paths to the DeepLesion3D data.
164
165    Args:
166        path: Filepath to a folder where the data is downloaded for further processing.
167        category: The lesion category. By default, the lesions of all categories are returned.
168        download: Whether to download the data if it is not present.
169
170    Returns:
171        List of filepaths for the image data.
172        List of filepaths for the label data.
173    """
174    data_dir = get_deeplesion_data(path, download)
175
176    with open(os.path.join(data_dir, "categories.json")) as f:
177        categories = json.load(f)
178
179    if category is not None and category not in CATEGORIES:
180        raise ValueError(f"'{category}' is not a valid category. Choose from {CATEGORIES}.")
181
182    fnames = natsorted(fname for fname, cat in categories.items() if category is None or cat == category)
183    raw_paths = [os.path.join(data_dir, "images", fname) for fname in fnames]
184    label_paths = [os.path.join(data_dir, "labels", fname) for fname in fnames]
185
186    assert len(raw_paths) == len(label_paths) and len(raw_paths) > 0
187    return raw_paths, label_paths
188
189
190def get_deeplesion_dataset(
191    path: Union[os.PathLike, str],
192    patch_shape: Tuple[int, ...],
193    category: Optional[Literal["Abdominal", "Bone", "Kidney", "Liver", "Lung", "Mediastinal", "Other"]] = None,
194    resize_inputs: bool = False,
195    download: bool = False,
196    **kwargs
197) -> Dataset:
198    """Get the DeepLesion3D dataset for lesion segmentation in CT.
199
200    Args:
201        path: Filepath to a folder where the data is downloaded for further processing.
202        patch_shape: The patch shape to use for training.
203        category: The lesion category. By default, the lesions of all categories are used.
204        resize_inputs: Whether to resize inputs to the desired patch shape.
205        download: Whether to download the data if it is not present.
206        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
207
208    Returns:
209        The segmentation dataset.
210    """
211    raw_paths, label_paths = get_deeplesion_paths(path, category, download)
212
213    if resize_inputs:
214        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
215        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
216            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
217        )
218
219    return torch_em.default_segmentation_dataset(
220        raw_paths=raw_paths,
221        raw_key="data",
222        label_paths=label_paths,
223        label_key="data",
224        patch_shape=patch_shape,
225        is_seg_dataset=True,
226        **kwargs
227    )
228
229
230def get_deeplesion_loader(
231    path: Union[os.PathLike, str],
232    batch_size: int,
233    patch_shape: Tuple[int, ...],
234    category: Optional[Literal["Abdominal", "Bone", "Kidney", "Liver", "Lung", "Mediastinal", "Other"]] = None,
235    resize_inputs: bool = False,
236    download: bool = False,
237    **kwargs
238) -> DataLoader:
239    """Get the DeepLesion3D dataloader for lesion segmentation in CT.
240
241    Args:
242        path: Filepath to a folder where the data is downloaded for further processing.
243        batch_size: The batch size for training.
244        patch_shape: The patch shape to use for training.
245        category: The lesion category. By default, the lesions of all categories are used.
246        resize_inputs: Whether to resize inputs to the desired patch shape.
247        download: Whether to download the data if it is not present.
248        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
249
250    Returns:
251        The DataLoader.
252    """
253    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
254    dataset = get_deeplesion_dataset(path, patch_shape, category, resize_inputs, download, **ds_kwargs)
255    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)