torch_em.data.datasets.medical.lnq

  1"""The LNQ dataset contains annotations for mediastinal lymph node segmentation in thorax CT scans.
  2
  3The dataset was curated for the LNQ2023 MICCAI challenge (https://lnq2023.grand-challenge.org).
  4It comprises 513 CT scans of cancer patients, each with a DICOM SEG object annotating the clinically
  5relevant mediastinal lymph nodes (larger than 1 cm). The annotations are partial for most cases,
  6i.e. not every lymph node in a scan is annotated. The data is hosted on TCIA as the
  7'Mediastinal-Lymph-Node-SEG' collection (https://doi.org/10.7937/QVAZ-JA09) under the CC BY 4.0 license.
  8
  9The DICOM series are converted to hdf5 files with the keys 'raw' (the CT in HU) and 'labels'.
 10The label ids are: 0 = background, 1 = lymph node. Note that each scan is annotated with a single
 11segment that covers all annotated lymph nodes, i.e. the individual nodes are not separated.
 12
 13The dataset is from the publication https://doi.org/10.59275/j.melba.2025-1gb5.
 14Please cite it if you use this dataset in your research.
 15
 16NOTE: The DICOM conversion requires 'pydicom'. Install it with 'pip install pydicom'.
 17"""
 18
 19import os
 20from glob import glob
 21from tqdm import tqdm
 22from natsort import natsorted
 23from typing import Union, Tuple, List
 24
 25import numpy as np
 26import pandas as pd
 27
 28from torch.utils.data import Dataset, DataLoader
 29
 30import torch_em
 31
 32from .. import util
 33
 34
 35URL = "https://www.cancerimagingarchive.net/wp-content/uploads/Mediastinal-Lymph-Node-SEG-DA-RAD.tcia"
 36
 37LABEL_IDS = {"background": 0, "lymph_node": 1}
 38
 39
 40def _load_ct_series(series_dir):
 41    import pydicom
 42
 43    slices = [pydicom.dcmread(p) for p in glob(os.path.join(series_dir, "*.dcm"))]
 44    slices = [s for s in slices if hasattr(s, "ImagePositionPatient")]
 45    slices.sort(key=lambda s: float(s.ImagePositionPatient[2]))
 46
 47    z_positions = np.array([float(s.ImagePositionPatient[2]) for s in slices])
 48    volume = np.stack([s.pixel_array.astype(np.float32) for s in slices])  # (Z, Y, X)
 49
 50    slope = float(getattr(slices[0], "RescaleSlope", 1.0))
 51    intercept = float(getattr(slices[0], "RescaleIntercept", 0.0))
 52    volume = (volume * slope + intercept).astype(np.int16)
 53
 54    return volume, z_positions
 55
 56
 57def _load_seg_series(series_dir, shape, z_positions):
 58    import pydicom
 59
 60    seg_path = glob(os.path.join(series_dir, "*.dcm"))[0]
 61    seg = pydicom.dcmread(seg_path)
 62
 63    frames = seg.pixel_array
 64    if frames.ndim == 2:  # A single frame is returned without the frame axis.
 65        frames = frames[None]
 66
 67    labels = np.zeros(shape, dtype=np.uint8)
 68    for frame, frame_info in zip(frames, seg.PerFrameFunctionalGroupsSequence):
 69        frame_z = float(frame_info.PlanePositionSequence[0].ImagePositionPatient[2])
 70
 71        # Match the frame to the CT slice with the closest z-position.
 72        z = int(np.argmin(np.abs(z_positions - frame_z)))
 73        assert abs(z_positions[z] - frame_z) < 1.0, "Could not match the segmentation frame to a CT slice."
 74        assert frame.shape == shape[1:], f"Frame shape {frame.shape} does not match the CT shape {shape[1:]}."
 75
 76        labels[z][frame > 0] = 1
 77
 78    return labels
 79
 80
 81def _preprocess_inputs(path, dicom_dir, csv_path):
 82    import h5py
 83
 84    preprocessed_dir = os.path.join(path, "preprocessed")
 85    os.makedirs(preprocessed_dir, exist_ok=True)
 86
 87    df = pd.read_csv(csv_path)
 88    for case_id, case_df in tqdm(df.groupby("Subject ID"), desc="Convert the LNQ DICOM series to hdf5"):
 89        volume_path = os.path.join(preprocessed_dir, f"{case_id}.h5")
 90        if os.path.exists(volume_path):
 91            continue
 92
 93        ct_uid = case_df[case_df["Modality"] == "CT"]["Series UID"].iloc[0]
 94        seg_uid = case_df[case_df["Modality"] == "SEG"]["Series UID"].iloc[0]
 95
 96        raw, z_positions = _load_ct_series(os.path.join(dicom_dir, ct_uid))
 97        labels = _load_seg_series(os.path.join(dicom_dir, seg_uid), raw.shape, z_positions)
 98
 99        # The file is written to a temporary path first, so that an interrupted run does not leave a corrupt file.
100        with h5py.File(f"{volume_path}.tmp", "w") as f:
101            f.create_dataset("raw", data=raw, compression="gzip")
102            f.create_dataset("labels", data=labels, compression="gzip")
103
104        os.rename(f"{volume_path}.tmp", volume_path)
105
106    return preprocessed_dir
107
108
109def get_lnq_data(path: Union[os.PathLike, str], download: bool = False) -> str:
110    """Download the LNQ dataset.
111
112    Args:
113        path: Filepath to a folder where the data is downloaded for further processing.
114        download: Whether to download the data if it is not present.
115
116    Returns:
117        Filepath where the data is downloaded and preprocessed.
118    """
119    preprocessed_dir = os.path.join(path, "preprocessed")
120    if os.path.exists(preprocessed_dir) and len(glob(os.path.join(preprocessed_dir, "*.h5"))) == 513:
121        return preprocessed_dir
122
123    os.makedirs(path, exist_ok=True)
124
125    dicom_dir = os.path.join(path, "dicom")
126    csv_path = os.path.join(path, "metadata")
127    if not os.path.exists(f"{csv_path}.csv"):
128        util.download_source_tcia(
129            path=os.path.join(path, "Mediastinal-Lymph-Node-SEG-DA-RAD.tcia"),
130            url=URL,
131            dst=dicom_dir,
132            csv_filename=csv_path,
133            download=download,
134        )
135
136    return _preprocess_inputs(path, dicom_dir, f"{csv_path}.csv")
137
138
139def get_lnq_paths(path: Union[os.PathLike, str], download: bool = False) -> Tuple[List[str], List[str]]:
140    """Get paths to the LNQ data.
141
142    Args:
143        path: Filepath to a folder where the data is downloaded for further processing.
144        download: Whether to download the data if it is not present.
145
146    Returns:
147        List of filepaths for the image data.
148        List of filepaths for the label data.
149    """
150    data_dir = get_lnq_data(path, download)
151    volume_paths = natsorted(glob(os.path.join(data_dir, "*.h5")))
152    return volume_paths, volume_paths
153
154
155def get_lnq_dataset(
156    path: Union[os.PathLike, str],
157    patch_shape: Tuple[int, ...],
158    resize_inputs: bool = False,
159    download: bool = False,
160    **kwargs
161) -> Dataset:
162    """Get the LNQ dataset for mediastinal lymph node segmentation.
163
164    Args:
165        path: Filepath to a folder where the data is downloaded for further processing.
166        patch_shape: The patch shape to use for training.
167        resize_inputs: Whether to resize inputs to the desired patch shape.
168        download: Whether to download the data if it is not present.
169        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
170
171    Returns:
172        The segmentation dataset.
173    """
174    raw_paths, label_paths = get_lnq_paths(path, download)
175
176    if resize_inputs:
177        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
178        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
179            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
180        )
181
182    return torch_em.default_segmentation_dataset(
183        raw_paths=raw_paths,
184        raw_key="raw",
185        label_paths=label_paths,
186        label_key="labels",
187        patch_shape=patch_shape,
188        is_seg_dataset=True,
189        **kwargs
190    )
191
192
193def get_lnq_loader(
194    path: Union[os.PathLike, str],
195    batch_size: int,
196    patch_shape: Tuple[int, ...],
197    resize_inputs: bool = False,
198    download: bool = False,
199    **kwargs
200) -> DataLoader:
201    """Get the LNQ dataloader for mediastinal lymph node segmentation.
202
203    Args:
204        path: Filepath to a folder where the data is downloaded for further processing.
205        batch_size: The batch size for training.
206        patch_shape: The patch shape to use for training.
207        resize_inputs: Whether to resize inputs to the desired patch shape.
208        download: Whether to download the data if it is not present.
209        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
210
211    Returns:
212        The DataLoader.
213    """
214    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
215    dataset = get_lnq_dataset(path, patch_shape, resize_inputs, download, **ds_kwargs)
216    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)