torch_em.data.datasets.medical.lndb

  1"""The LNDb dataset contains annotations for lung nodule segmentation in chest CT.
  2
  3It consists of 294 chest CT that were collected at the Centro Hospitalar e Universitario de Sao Joao in Porto
  4between 2016 and 2018. The findings were marked by up to three radiologists out of a group of five, who segmented
  5every finding that they considered a nodule with a diameter of at least 3 mm. Non-nodules and nodules below 3 mm
  6were only marked with a centroid and are not segmented. This module uses the 236 CT of the challenge training set,
  7which are the only scans with public segmentations. The 58 test CT are only released in the newer version of the
  8dataset (https://doi.org/10.5281/zenodo.7153205) and come with nodule centroids, but without segmentations, so
  9they are not included here.
 10
 11The CT are distributed as MetaImage ('mhd' / 'raw') volumes and the annotations as one mask per radiologist, in
 12which the voxel values are the finding ids of that radiologist. This module joins the per-radiologist findings
 13into nodules with the official matching of 'trainNodules_gt.csv' (findings of different radiologists are the same
 14nodule if their centroids are closer than the maximum of their radii and 3 mm) and stores the scan and the nodule
 15masks together in hdf5 files. Each nodule gets one instance id and the files contain the nodule masks for all
 16consensus levels ('labels/consensus_1' to 'labels/consensus_3'), where consensus level n contains the voxels that
 17at least n radiologists marked as part of the nodule (level 1 is the union of all readings). The number of
 18radiologists that annotated each voxel is stored in 'labels/n_readers' and the number of radiologists that read
 19the scan in the 'n_readers' attribute (59 scans were read by one, 110 by two and 67 by three radiologists,
 20so the higher consensus levels are empty for the scans with fewer readers). Findings that all readers
 21considered non-nodules and findings that were not segmented are not included, while the findings below the
 22equivalent diameter of 3 mm that the official example scripts filter out are kept.
 23
 24The volumes are converted to the (Z, Y, X) axis order, which is the order the MetaImage files are stored in.
 25
 26The data is located at https://doi.org/10.5281/zenodo.6613714. NOTE: The license field of the Zenodo record
 27states CC BY 4.0, but its description states CC BY-NC-ND 4.0, so please make sure that you are allowed to use
 28the data for your purpose.
 29
 30This dataset is from the publication https://doi.org/10.48550/arXiv.1911.08434 and the challenge it was used for
 31is described in https://doi.org/10.1016/j.media.2021.102027.
 32Please cite them if you use this dataset in your research.
 33"""
 34
 35import os
 36import csv
 37from glob import glob
 38from tqdm import tqdm
 39from natsort import natsorted
 40from collections import defaultdict
 41from typing import Union, Tuple, List
 42
 43import numpy as np
 44
 45from torch.utils.data import Dataset, DataLoader
 46
 47import torch_em
 48
 49from .. import util
 50
 51
 52URLS = {
 53    "data0": "https://zenodo.org/records/6613714/files/data0.rar",
 54    "data1": "https://zenodo.org/records/6613714/files/data1.rar",
 55    "data2": "https://zenodo.org/records/6613714/files/data2.rar",
 56    "data3": "https://zenodo.org/records/6613714/files/data3.rar",
 57    "data4": "https://zenodo.org/records/6613714/files/data4.rar",
 58    "data5": "https://zenodo.org/records/6613714/files/data5.rar",
 59    "masks": "https://zenodo.org/records/6613714/files/masks.rar",
 60    "csvs": "https://zenodo.org/records/6613714/files/trainset_csv.zip",
 61}
 62
 63CHECKSUMS = {
 64    "data0": "2ae6b7b760f5286bde2f51b32fcdedbf4e1e35aee575cb971bbde5d1aef517d7",
 65    "data1": "f5b33500127e09c0c320561c9c72e72a7934f02214bb091d717d95ed27977057",
 66    "data2": "a8819b5ad36023cb85d1bf840a136c413bbac38dbd3cbc2b4bddb2ef43acf0cf",
 67    "data3": "05a4b817c112aec650e78e6ba6908737a26550510de5a2b62d2b8c0dbc3059e7",
 68    "data4": "ffa1a2de6525a1409c7d5a0d0bbdf9ce758219cca7914c08471adf3cb4405d15",
 69    "data5": "6e94ebcc2650879feb6928b30f8bd6f8675a1c0557d9e7c06852fba81e5cea6d",
 70    "masks": "cdc9cfbbb868ec6fc6e2612afc9bc12cd6e07f8d345610a4dffbf76e1bcc7980",
 71    "csvs": "cda2e6377344dd90b7926f2c10da29d5260ca59337bd8f1c1490639ffa8fbb94",
 72}
 73
 74# The nodule masks are instance labels, so the label ids are the nodule ids of a scan.
 75LABEL_IDS = {"background": 0, "nodule": "1, 2, ... (one id per nodule)"}
 76
 77MAX_CONSENSUS_LEVEL = 3
 78
 79N_VOLUMES = 236
 80
 81N_MASKS = 480
 82
 83METAIMAGE_DTYPES = {
 84    "MET_CHAR": "int8",
 85    "MET_UCHAR": "uint8",
 86    "MET_SHORT": "int16",
 87    "MET_USHORT": "uint16",
 88    "MET_INT": "int32",
 89    "MET_UINT": "uint32",
 90    "MET_FLOAT": "float32",
 91    "MET_DOUBLE": "float64",
 92}
 93
 94
 95def _read_metaimage(mhd_path):
 96    """Read a MetaImage volume with the axis order (Z, Y, X).
 97
 98    The header is a plain text file that refers to the raw data of the volume, so no additional dependency
 99    (such as SimpleITK) is needed to read it.
100    """
101    header = {}
102    with open(mhd_path, "r") as f:
103        for line in f:
104            if "=" not in line:
105                continue
106            key, value = line.split("=", 1)
107            header[key.strip()] = value.strip()
108
109    shape = [int(v) for v in header["DimSize"].split()][::-1]
110    dtype = METAIMAGE_DTYPES[header["ElementType"]]
111    raw_path = os.path.join(os.path.dirname(mhd_path), os.path.basename(header["ElementDataFile"]))
112    return np.fromfile(raw_path, dtype=dtype).reshape(shape)
113
114
115def _read_csv(csv_path):
116    with open(csv_path, "r") as f:
117        return list(csv.DictReader(f))
118
119
120def _get_findings_per_scan(csv_dir):
121    """Group the joined findings of 'trainNodules_gt.csv' by scan.
122
123    Each finding is a list of the (radiologist id, finding id of that radiologist) pairs that were joined into it.
124    The findings that all readers considered non-nodules are skipped.
125    """
126    findings = defaultdict(list)
127    for row in _read_csv(os.path.join(csv_dir, "trainNodules_gt.csv")):
128        if int(row["Nodule"]) != 1:
129            continue
130        rad_ids = [int(v) for v in row["RadID"].split(",")]
131        rad_finding_ids = [int(v) for v in row["RadFindingID"].split(",")]
132        findings[int(row["LNDbID"])].append(list(zip(rad_ids, rad_finding_ids)))
133    return findings
134
135
136def _build_nodule_labels(mask_dir, scan_id, findings, shape):
137    """Derive the instance labels for all consensus levels from the per-radiologist masks of one scan."""
138    masks = {}
139    for mask_path in glob(os.path.join(mask_dir, f"LNDb-{scan_id:04d}_rad*.mhd")):
140        rad_id = int(os.path.basename(mask_path).split("_rad")[1][:-len(".mhd")])
141        masks[rad_id] = _read_metaimage(mask_path)
142
143    n_readers = np.zeros(shape, dtype="uint8")
144    consensus = {level: np.zeros(shape, dtype="uint16") for level in range(1, MAX_CONSENSUS_LEVEL + 1)}
145
146    nodule_id = 0
147    for finding in findings:
148        # Count how many radiologists annotated each voxel as part of this nodule.
149        counts = np.zeros(shape, dtype="uint8")
150        for rad_id, rad_finding_id in finding:
151            if rad_id in masks:
152                counts += (masks[rad_id] == rad_finding_id)
153
154        # Findings that were only marked with a centroid are not segmented and are skipped.
155        if not counts.any():
156            continue
157
158        nodule_id += 1
159        n_readers = np.maximum(n_readers, counts)
160        for level, labels in consensus.items():
161            labels[counts >= level] = nodule_id
162
163    return consensus, n_readers, len(masks)
164
165
166def _preprocess_lndb(data_dir, mask_dir, csv_dir, preprocessed_dir):
167    import h5py
168
169    findings_per_scan = _get_findings_per_scan(csv_dir)
170    scan_ids = [int(row["LNDbID"]) for row in _read_csv(os.path.join(csv_dir, "trainCTs.csv"))]
171    os.makedirs(preprocessed_dir, exist_ok=True)
172
173    for scan_id in tqdm(scan_ids, desc="Preprocessing the LNDb scans"):
174        volume_path = os.path.join(preprocessed_dir, f"LNDb-{scan_id:04d}.h5")
175        if os.path.exists(volume_path):
176            continue
177
178        raw = _read_metaimage(os.path.join(data_dir, f"LNDb-{scan_id:04d}.mhd"))
179        consensus, n_readers, n_masks = _build_nodule_labels(
180            mask_dir, scan_id, findings_per_scan[scan_id], raw.shape
181        )
182
183        # The file is written to a temporary path first, so that an interrupted run leaves no corrupt file.
184        with h5py.File(f"{volume_path}.tmp", "w") as f:
185            f.attrs["scan_id"] = scan_id
186            f.attrs["n_readers"] = n_masks
187            f.create_dataset("raw", data=raw, compression="gzip")
188            f.create_dataset("labels/n_readers", data=n_readers, compression="gzip")
189            for level, labels in consensus.items():
190                f.create_dataset(f"labels/consensus_{level}", data=labels, compression="gzip")
191
192        os.rename(f"{volume_path}.tmp", volume_path)
193
194
195def get_lndb_data(path: Union[os.PathLike, str], download: bool = False) -> str:
196    """Download the LNDb dataset.
197
198    Args:
199        path: Filepath to a folder where the data is downloaded for further processing.
200        download: Whether to download the data if it is not present.
201
202    Returns:
203        Filepath where the preprocessed data is stored.
204    """
205    preprocessed_dir = os.path.join(path, "preprocessed")
206    if len(glob(os.path.join(preprocessed_dir, "*.h5"))) == N_VOLUMES:
207        return preprocessed_dir
208
209    os.makedirs(path, exist_ok=True)
210
211    # The scans are spread over six rar archives, which all unpack into the same folder. A marker file is written
212    # for every archive that was unpacked, so that an interrupted run does not download the 24 GB again.
213    data_dir = os.path.join(path, "data")
214    for name in [f"data{i}" for i in range(6)]:
215        marker_path = os.path.join(path, f"{name}.extracted")
216        if os.path.exists(marker_path):
217            continue
218        rar_path = os.path.join(path, f"{name}.rar")
219        util.download_source(path=rar_path, url=URLS[name], download=download, checksum=CHECKSUMS[name])
220        util.unzip_rarfile(rar_path=rar_path, dst=data_dir)
221        open(marker_path, "w").close()
222
223    # There is one mask per radiologist and scan, which gives 480 masks for the 236 scans.
224    mask_dir = os.path.join(path, "masks")
225    if len(glob(os.path.join(mask_dir, "*.mhd"))) != N_MASKS:
226        rar_path = os.path.join(path, "masks.rar")
227        util.download_source(path=rar_path, url=URLS["masks"], download=download, checksum=CHECKSUMS["masks"])
228        util.unzip_rarfile(rar_path=rar_path, dst=path)
229
230    csv_dir = os.path.join(path, "trainset_csv")
231    if not os.path.exists(os.path.join(csv_dir, "trainNodules_gt.csv")):
232        zip_path = os.path.join(path, "trainset_csv.zip")
233        util.download_source(path=zip_path, url=URLS["csvs"], download=download, checksum=CHECKSUMS["csvs"])
234        util.unzip(zip_path=zip_path, dst=csv_dir)
235
236    _preprocess_lndb(data_dir, mask_dir, csv_dir, preprocessed_dir)
237    return preprocessed_dir
238
239
240def get_lndb_paths(path: Union[os.PathLike, str], download: bool = False) -> List[str]:
241    """Get paths to the LNDb data.
242
243    Args:
244        path: Filepath to a folder where the data is downloaded for further processing.
245        download: Whether to download the data if it is not present.
246
247    Returns:
248        List of filepaths for the hdf5 files, which contain the image data ('raw') and the label data
249        ('labels/consensus_<level>' and 'labels/n_readers').
250    """
251    data_dir = get_lndb_data(path, download)
252    volume_paths = natsorted(glob(os.path.join(data_dir, "*.h5")))
253    assert len(volume_paths) > 0
254
255    return volume_paths
256
257
258def get_lndb_dataset(
259    path: Union[os.PathLike, str],
260    patch_shape: Tuple[int, ...],
261    consensus_level: int = 1,
262    resize_inputs: bool = False,
263    download: bool = False,
264    **kwargs
265) -> Dataset:
266    """Get the LNDb dataset for lung nodule segmentation.
267
268    Args:
269        path: Filepath to a folder where the data is downloaded for further processing.
270        patch_shape: The patch shape to use for training.
271        consensus_level: The minimum number of radiologists (1 to 3) that have to agree on a voxel for it to be part
272            of a nodule. 1 corresponds to the union of all radiologist annotations.
273        resize_inputs: Whether to resize inputs to the desired patch shape.
274        download: Whether to download the data if it is not present.
275        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
276
277    Returns:
278        The segmentation dataset.
279    """
280    if consensus_level not in range(1, MAX_CONSENSUS_LEVEL + 1):
281        raise ValueError(f"'{consensus_level}' is not a valid consensus level. Please choose a value from 1 to 3.")
282
283    volume_paths = get_lndb_paths(path, download)
284
285    if resize_inputs:
286        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
287        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
288            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
289        )
290
291    return torch_em.default_segmentation_dataset(
292        raw_paths=volume_paths,
293        raw_key="raw",
294        label_paths=volume_paths,
295        label_key=f"labels/consensus_{consensus_level}",
296        patch_shape=patch_shape,
297        is_seg_dataset=True,
298        **kwargs
299    )
300
301
302def get_lndb_loader(
303    path: Union[os.PathLike, str],
304    batch_size: int,
305    patch_shape: Tuple[int, ...],
306    consensus_level: int = 1,
307    resize_inputs: bool = False,
308    download: bool = False,
309    **kwargs
310) -> DataLoader:
311    """Get the LNDb dataloader for lung nodule segmentation.
312
313    Args:
314        path: Filepath to a folder where the data is downloaded for further processing.
315        batch_size: The batch size for training.
316        patch_shape: The patch shape to use for training.
317        consensus_level: The minimum number of radiologists (1 to 3) that have to agree on a voxel for it to be part
318            of a nodule. 1 corresponds to the union of all radiologist annotations.
319        resize_inputs: Whether to resize inputs to the desired patch shape.
320        download: Whether to download the data if it is not present.
321        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
322
323    Returns:
324        The DataLoader.
325    """
326    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
327    dataset = get_lndb_dataset(path, patch_shape, consensus_level, resize_inputs, download, **ds_kwargs)
328    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)