torch_em.data.datasets.medical.lidc_idri

  1"""The LIDC-IDRI dataset contains annotations for lung nodule segmentation in thoracic CT.
  2
  3It consists of 1018 CT scans of 1010 patients, in which up to four thoracic radiologists outlined the nodules
  4with a diameter of at least 3 mm. The CT scans are distributed as DICOM series (the CR / DX scans of the collection
  5are skipped) and the annotations as XML files with per-radiologist nodule contours. This module stacks the DICOM
  6series into volumes, rasterizes the contours onto the CT grid and stores them together in hdf5 files.
  7
  8The contours of a radiologist are rasterized per slice with `skimage.draw.polygon`. Following the LIDC convention
  9(and pylidc), the contour points themselves are not part of the nodule and 'exclusion' contours are subtracted.
 10The individual radiologist annotations are grouped into nodules by voxel overlap: annotations of different
 11radiologists that overlap are treated as the same nodule. Each nodule gets one instance id and the hdf5 files
 12store the nodule masks for all consensus levels ('labels/consensus_1' to 'labels/consensus_4'), where consensus
 13level n contains the voxels that at least n radiologists marked as part of the nodule (level 1 is the union,
 14level 4 is the intersection of all four readings). The number of radiologists that annotated each voxel is stored
 15in 'labels/n_readers'. Nodules < 3 mm and non-nodules (single point marks) are not rasterized.
 16
 17NOTE: This requires the pydicom python package.
 18
 19The dataset is located at https://www.cancerimagingarchive.net/collection/lidc-idri/.
 20
 21This dataset is from the publication https://doi.org/10.1118/1.3528204.
 22The data was released at https://doi.org/10.7937/K9/TCIA.2015.LO9QL9SX.
 23Please cite it if you use this dataset in your research.
 24"""
 25
 26import os
 27import json
 28from glob import glob
 29from tqdm import tqdm
 30from natsort import natsorted
 31from collections import defaultdict
 32from typing import Union, Tuple, List
 33import xml.etree.ElementTree as ET
 34
 35import numpy as np
 36import requests
 37
 38from torch.utils.data import Dataset, DataLoader
 39
 40import torch_em
 41
 42from .. import util
 43
 44
 45URLS = {
 46    "images": f"{util.NBIA_API_URL}getSeries?Collection=LIDC-IDRI&Modality=CT",
 47    "annotations": "https://www.cancerimagingarchive.net/wp-content/uploads/LIDC-XML-only.zip",
 48}
 49
 50CHECKSUMS = {
 51    "images": None,  # The DICOM series are downloaded individually from TCIA.
 52    "annotations": "644557a3aa305602609c718b0cae33a93be762e8901a80bcd9df735b0ec6ab90",
 53}
 54
 55XML_NAMESPACE = "{http://www.nih.gov}"
 56
 57
 58def _load_dicom_volume(series_dir):
 59    """Stack a DICOM series into a volume with axes (z, y, x) and slices sorted by ascending patient z position.
 60
 61    Returns the volume in Hounsfield units, the SOP instance UID and the z position of each slice.
 62    """
 63    import pydicom
 64
 65    slices = [pydicom.dcmread(dcm_path) for dcm_path in natsorted(glob(os.path.join(series_dir, "*.dcm")))]
 66    slices = [dcm for dcm in slices if hasattr(dcm, "ImagePositionPatient")]
 67    slices.sort(key=lambda dcm: float(dcm.ImagePositionPatient[2]))
 68
 69    volume = np.stack([dcm.pixel_array for dcm in slices]).astype("float32")
 70    volume = volume * float(slices[0].RescaleSlope) + float(slices[0].RescaleIntercept)
 71    volume = np.round(volume).astype("int16")
 72
 73    sop_uids = [str(dcm.SOPInstanceUID) for dcm in slices]
 74    z_positions = np.array([float(dcm.ImagePositionPatient[2]) for dcm in slices])
 75    return volume, sop_uids, z_positions
 76
 77
 78def _parse_annotations(xml_path):
 79    """Parse the nodule (>= 3 mm) contours of all reading sessions in a LIDC XML file.
 80
 81    Returns a list with one entry per radiologist (reading session). Each entry is a list of nodules and each nodule
 82    is a list of contours (imageSOP_UID, imageZposition, inclusion, points), where points is an array of
 83    (row, column) coordinates. Nodules < 3 mm (a single point mark) are skipped.
 84    """
 85    root = ET.parse(xml_path).getroot()
 86    sessions = []
 87    for session in root.iter(f"{XML_NAMESPACE}readingSession"):
 88        nodules = []
 89        for nodule in session.iter(f"{XML_NAMESPACE}unblindedReadNodule"):
 90            contours = []
 91            for roi in nodule.iter(f"{XML_NAMESPACE}roi"):
 92                points = np.array([
 93                    [int(edge.find(f"{XML_NAMESPACE}yCoord").text), int(edge.find(f"{XML_NAMESPACE}xCoord").text)]
 94                    for edge in roi.iter(f"{XML_NAMESPACE}edgeMap")
 95                ])
 96                inclusion = roi.find(f"{XML_NAMESPACE}inclusion").text.strip().upper() == "TRUE"
 97                sop_uid = roi.find(f"{XML_NAMESPACE}imageSOP_UID").text.strip()
 98                z_position = float(roi.find(f"{XML_NAMESPACE}imageZposition").text)
 99                contours.append((sop_uid, z_position, inclusion, points))
100            # Nodules < 3 mm are marked with a single point and are not segmented.
101            if all(len(points) < 3 for _, _, _, points in contours):
102                continue
103            nodules.append(contours)
104        sessions.append(nodules)
105    return sessions
106
107
108def _find_slice(sop_uid, z_position, slice_ids, z_positions):
109    """Find the slice a contour belongs to via its SOP instance UID.
110
111    For a few scans the XML files reference SOP instance UIDs that do not exist in the DICOM series
112    (e.g. LIDC-IDRI-0017). In this case the slice is found via the z position, which has to match a slice
113    position exactly.
114    """
115    if sop_uid in slice_ids:
116        return slice_ids[sop_uid]
117    z = int(np.argmin(np.abs(z_positions - z_position)))
118    if abs(z_positions[z] - z_position) > 1e-3:
119        raise ValueError(f"Cannot find the slice for the SOP instance UID {sop_uid} at z position {z_position}.")
120    return z
121
122
123def _rasterize_nodule(contours, shape, slice_ids, z_positions):
124    """Rasterize the contours of one nodule annotation onto the CT grid."""
125    from skimage.draw import polygon
126
127    mask = np.zeros(shape, dtype="bool")
128    # Fill the inclusion contours first, then subtract the exclusion contours.
129    for inclusion in (True, False):
130        for sop_uid, z_position, is_inclusion, points in contours:
131            if is_inclusion != inclusion or len(points) < 3:
132                continue
133            z = _find_slice(sop_uid, z_position, slice_ids, z_positions)
134            rr, cc = polygon(points[:, 0], points[:, 1], shape=shape[1:])
135            if inclusion:
136                mask[z, rr, cc] = True
137            else:
138                mask[z, rr, cc] = False
139            # The contour points are not part of the nodule (they lie just outside of it).
140            mask[z, points[:, 0], points[:, 1]] = False
141    return mask
142
143
144def _build_nodule_labels(sessions, shape, slice_ids, z_positions):
145    """Group the per-radiologist nodule annotations into nodules and derive the consensus labels.
146
147    Returns the instance labels for the consensus levels 1 to 4 and the number of readers per voxel.
148    """
149    # Rasterize all annotations and keep only the crop around each annotation to save memory.
150    annotations = []
151    for reader_id, nodules in enumerate(sessions):
152        for contours in nodules:
153            mask = _rasterize_nodule(contours, shape, slice_ids, z_positions)
154            if not mask.any():
155                continue
156            bbox = tuple(slice(int(c.min()), int(c.max()) + 1) for c in np.where(mask))
157            annotations.append((reader_id, bbox, mask[bbox]))
158
159    def intersect(bbox_a, bbox_b):
160        bbox = tuple(slice(max(a.start, b.start), min(a.stop, b.stop)) for a, b in zip(bbox_a, bbox_b))
161        return bbox if all(b.stop > b.start for b in bbox) else None
162
163    def crop(mask, bbox, sub_bbox):
164        return mask[tuple(slice(s.start - b.start, s.stop - b.start) for s, b in zip(sub_bbox, bbox))]
165
166    # Group the annotations by voxel overlap (union-find over the pairwise overlaps).
167    parents = list(range(len(annotations)))
168
169    def find(i):
170        while parents[i] != i:
171            parents[i] = parents[parents[i]]
172            i = parents[i]
173        return i
174
175    for i in range(len(annotations)):
176        for j in range(i + 1, len(annotations)):
177            _, bbox_i, mask_i = annotations[i]
178            _, bbox_j, mask_j = annotations[j]
179            overlap = intersect(bbox_i, bbox_j)
180            if overlap is not None and np.any(crop(mask_i, bbox_i, overlap) & crop(mask_j, bbox_j, overlap)):
181                parents[find(i)] = find(j)
182
183    groups = defaultdict(list)
184    for i in range(len(annotations)):
185        groups[find(i)].append(i)
186
187    n_readers = np.zeros(shape, dtype="uint8")
188    consensus = {level: np.zeros(shape, dtype="uint16") for level in range(1, 5)}
189    for nodule_id, members in enumerate(groups.values(), start=1):
190        # Count the readers per voxel (a reader may have annotated the same nodule more than once).
191        counts = np.zeros(shape, dtype="uint8")
192        for reader_id in set(annotations[i][0] for i in members):
193            reader_mask = np.zeros(shape, dtype="bool")
194            for i in members:
195                if annotations[i][0] == reader_id:
196                    reader_mask[annotations[i][1]] |= annotations[i][2]
197            counts += reader_mask
198        n_readers = np.maximum(n_readers, counts)
199        for level in consensus:
200            consensus[level][counts >= level] = nodule_id
201
202    return consensus, n_readers
203
204
205def _preprocess_lidc_idri(dicom_dir, xml_dir, series_metadata, preprocessed_dir):
206    import h5py
207
208    # Map the series UIDs to the XML files. A few series have multiple (identical) XML files
209    # and the file for LIDC-IDRI-0101 was resubmitted with a correction, which takes precedence.
210    xml_paths = {}
211    all_xml_paths = glob(os.path.join(xml_dir, "**", "*.xml"), recursive=True)
212    all_xml_paths += glob(os.path.join(xml_dir, "..", "*.xml"))
213    for xml_path in natsorted(all_xml_paths):
214        with open(xml_path, "r", errors="replace") as f:
215            header = f.read(4096)
216        if "LidcReadMessage" not in header:  # Skip the CXR annotations.
217            continue
218        start = header.index("<SeriesInstanceUid>") + len("<SeriesInstanceUid>")
219        series_uid = header[start:header.index("</SeriesInstanceUid>")].strip()
220        if series_uid not in xml_paths or "correction" in os.path.basename(xml_path):
221            xml_paths[series_uid] = xml_path
222
223    # The 8 patients with two CT series get one file per series (suffixed with the series number).
224    series_per_patient = defaultdict(list)
225    for series in series_metadata:
226        series_per_patient[series["PatientID"]].append(series)
227
228    os.makedirs(preprocessed_dir, exist_ok=True)
229    for series in tqdm(series_metadata, desc="Preprocess LIDC-IDRI"):
230        patient_id, series_uid = series["PatientID"], series["SeriesInstanceUID"]
231        name = patient_id if len(series_per_patient[patient_id]) == 1 else f"{patient_id}_{series['SeriesNumber']}"
232        out_path = os.path.join(preprocessed_dir, f"{name}.h5")
233        if os.path.exists(out_path):
234            continue
235
236        series_dir = os.path.join(dicom_dir, series_uid)
237        if not os.path.exists(series_dir):
238            continue
239
240        volume, sop_uids, z_positions = _load_dicom_volume(series_dir)
241        slice_ids = {sop_uid: z for z, sop_uid in enumerate(sop_uids)}
242        sessions = _parse_annotations(xml_paths[series_uid])
243        consensus, n_readers = _build_nodule_labels(sessions, volume.shape, slice_ids, z_positions)
244
245        tmp_path = out_path + ".tmp"
246        with h5py.File(tmp_path, "w") as f:
247            f.attrs["patient_id"] = patient_id
248            f.attrs["series_uid"] = series_uid
249            f.attrs["n_readers"] = len(sessions)
250            f.create_dataset("raw", data=volume, compression="gzip")
251            f.create_dataset("labels/n_readers", data=n_readers, compression="gzip")
252            for level, labels in consensus.items():
253                f.create_dataset(f"labels/consensus_{level}", data=labels, compression="gzip")
254        os.rename(tmp_path, out_path)
255
256
257def _get_series_metadata(path, download):
258    """Get the metadata of all CT series in the collection from the NBIA REST API."""
259    metadata_path = os.path.join(path, "lidc_idri_ct_series.json")
260    if not os.path.exists(metadata_path):
261        if not download:
262            raise RuntimeError(f"Cannot find the data at {path}, but download was set to False.")
263        response = requests.get(URLS["images"])
264        response.raise_for_status()
265        with open(metadata_path, "w") as f:
266            json.dump(response.json(), f, indent=2)
267
268    with open(metadata_path, "r") as f:
269        series_metadata = json.load(f)
270    return sorted(series_metadata, key=lambda series: (series["PatientID"], series["SeriesInstanceUID"]))
271
272
273def get_lidc_idri_data(path: Union[os.PathLike, str], download: bool = False) -> str:
274    """Download the LIDC-IDRI dataset.
275
276    The download is resumable: series that were already downloaded and volumes that were already converted are
277    skipped when the function is called again.
278
279    Args:
280        path: Filepath to a folder where the data is downloaded for further processing.
281        download: Whether to download the data if it is not present.
282
283    Returns:
284        Filepath where the preprocessed data is stored.
285    """
286    os.makedirs(path, exist_ok=True)
287    preprocessed_dir = os.path.join(path, "preprocessed")
288
289    # The list of CT series (the CR / DX scans in the collection are skipped).
290    series_metadata = _get_series_metadata(path, download)
291    if len(glob(os.path.join(preprocessed_dir, "*.h5"))) == len(series_metadata):
292        return preprocessed_dir
293
294    # Download the XML annotations.
295    xml_dir = os.path.join(path, "tcia-lidc-xml")
296    if not os.path.exists(xml_dir):
297        zip_path = os.path.join(path, "LIDC-XML-only.zip")
298        util.download_source(
299            path=zip_path, url=URLS["annotations"], download=download, checksum=CHECKSUMS["annotations"]
300        )
301        util.unzip(zip_path=zip_path, dst=path)
302
303    # Download the CT series from TCIA (series that were downloaded already are skipped).
304    dicom_dir = os.path.join(path, "dicom")
305    series_uids = [series["SeriesInstanceUID"] for series in series_metadata]
306    if download:
307        util.download_tcia_series(series_uids, dst=dicom_dir, csv_filename=os.path.join(path, "lidc_idri_series"))
308    elif not all(os.path.exists(os.path.join(dicom_dir, uid)) for uid in series_uids):
309        raise RuntimeError(f"Cannot find the data at {path}, but download was set to False.")
310
311    _preprocess_lidc_idri(dicom_dir, xml_dir, series_metadata, preprocessed_dir)
312    return preprocessed_dir
313
314
315def get_lidc_idri_paths(path: Union[os.PathLike, str], download: bool = False) -> List[str]:
316    """Get paths to the LIDC-IDRI data.
317
318    Args:
319        path: Filepath to a folder where the data is downloaded for further processing.
320        download: Whether to download the data if it is not present.
321
322    Returns:
323        List of filepaths for the hdf5 files, which contain the image data ('raw') and the label data
324        ('labels/consensus_<level>' and 'labels/n_readers').
325    """
326    data_dir = get_lidc_idri_data(path, download)
327    volume_paths = natsorted(glob(os.path.join(data_dir, "*.h5")))
328    return volume_paths
329
330
331def get_lidc_idri_dataset(
332    path: Union[os.PathLike, str],
333    patch_shape: Tuple[int, ...],
334    consensus_level: int = 1,
335    resize_inputs: bool = False,
336    download: bool = False,
337    **kwargs
338) -> Dataset:
339    """Get the LIDC-IDRI dataset for lung nodule segmentation.
340
341    Args:
342        path: Filepath to a folder where the data is downloaded for further processing.
343        patch_shape: The patch shape to use for training.
344        consensus_level: The minimum number of radiologists (1 to 4) that have to agree on a voxel for it to be part
345            of a nodule. 1 corresponds to the union and 4 to the intersection of all radiologist annotations.
346        resize_inputs: Whether to resize inputs to the desired patch shape.
347        download: Whether to download the data if it is not present.
348        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
349
350    Returns:
351        The segmentation dataset.
352    """
353    assert consensus_level in (1, 2, 3, 4), f"Invalid consensus level: {consensus_level}."
354    volume_paths = get_lidc_idri_paths(path, download)
355
356    if resize_inputs:
357        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
358        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
359            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
360        )
361
362    return torch_em.default_segmentation_dataset(
363        raw_paths=volume_paths,
364        raw_key="raw",
365        label_paths=volume_paths,
366        label_key=f"labels/consensus_{consensus_level}",
367        patch_shape=patch_shape,
368        is_seg_dataset=True,
369        **kwargs
370    )
371
372
373def get_lidc_idri_loader(
374    path: Union[os.PathLike, str],
375    batch_size: int,
376    patch_shape: Tuple[int, ...],
377    consensus_level: int = 1,
378    resize_inputs: bool = False,
379    download: bool = False,
380    **kwargs
381) -> DataLoader:
382    """Get the LIDC-IDRI dataloader for lung nodule segmentation.
383
384    Args:
385        path: Filepath to a folder where the data is downloaded for further processing.
386        batch_size: The batch size for training.
387        patch_shape: The patch shape to use for training.
388        consensus_level: The minimum number of radiologists (1 to 4) that have to agree on a voxel for it to be part
389            of a nodule. 1 corresponds to the union and 4 to the intersection of all radiologist annotations.
390        resize_inputs: Whether to resize inputs to the desired patch shape.
391        download: Whether to download the data if it is not present.
392        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
393
394    Returns:
395        The DataLoader.
396    """
397    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
398    dataset = get_lidc_idri_dataset(path, patch_shape, consensus_level, resize_inputs, download, **ds_kwargs)
399    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)