torch_em.data.datasets.medical.adrenal_acc

  1"""The Adrenal-ACC-Ki67-Seg dataset contains annotations for adrenocortical carcinoma segmentation
  2in contrast-enhanced abdominal CT.
  3
  4It consists of 53 CT volumes (one segmented CT series per patient, the collection also contains 71 additional
  5unsegmented CT series of the same patients) with binary tumor labels (1: adrenal tumor). The CT scans are
  6distributed as DICOM series and the labels as DICOM-SEG objects, which are converted and stored in hdf5 files
  7by this module.
  8
  9NOTE: This requires the pydicom python package.
 10
 11The dataset is located at https://www.cancerimagingarchive.net/collection/adrenal-acc-ki67-seg/.
 12
 13This dataset is from the publication https://doi.org/10.1016/j.crad.2020.01.012.
 14The data was released at https://doi.org/10.7937/1FPG-VM46.
 15Please cite it if you use this dataset in your research.
 16"""
 17
 18import os
 19import csv
 20from glob import glob
 21from tqdm import tqdm
 22from natsort import natsorted
 23from typing import Union, Tuple, List
 24
 25import numpy as np
 26
 27from torch.utils.data import Dataset, DataLoader
 28
 29import torch_em
 30
 31from .. import util
 32
 33
 34URL = "https://www.cancerimagingarchive.net/wp-content/uploads/Adrenal-ACC-Ki67-Seg_v1.tcia"
 35
 36# The DICOM series are downloaded individually from TCIA.
 37CHECKSUM = None
 38
 39LABEL_IDS = {"tumor": 1}
 40
 41
 42def _load_dicom_volume(series_dir):
 43    """Stack a DICOM series into a volume with axes (z, y, x) and slices sorted along the slice normal.
 44
 45    Returns the volume in Hounsfield units and the affine matrix that maps voxel indices (z, y, x)
 46    to DICOM patient coordinates.
 47    """
 48    import pydicom
 49
 50    slices = [pydicom.dcmread(dcm_path) for dcm_path in natsorted(glob(os.path.join(series_dir, "*.dcm")))]
 51    orientation = np.array([float(v) for v in slices[0].ImageOrientationPatient])
 52    row_dir, col_dir = orientation[:3], orientation[3:]
 53    normal = np.cross(row_dir, col_dir)
 54    slices.sort(key=lambda dcm: np.dot([float(v) for v in dcm.ImagePositionPatient], normal))
 55
 56    volume = np.stack([dcm.pixel_array for dcm in slices]).astype("float32")
 57    volume = volume * float(slices[0].RescaleSlope) + float(slices[0].RescaleIntercept)
 58    volume = np.round(volume).astype("int16")
 59
 60    positions = np.array([[float(v) for v in dcm.ImagePositionPatient] for dcm in slices])
 61    spacing = [float(v) for v in slices[0].PixelSpacing]  # The spacing between rows and between columns.
 62    affine = np.eye(4)
 63    affine[:3, 0] = (positions[-1] - positions[0]) / (len(slices) - 1)
 64    affine[:3, 1] = col_dir * spacing[0]
 65    affine[:3, 2] = row_dir * spacing[1]
 66    affine[:3, 3] = positions[0]
 67    return volume, affine
 68
 69
 70def _load_dicom_seg(seg_path):
 71    """Load a DICOM-SEG object as a label volume with axes (z, y, x), where the segment number is used as label id.
 72
 73    Returns the label volume and the affine matrix that maps its voxel indices to DICOM patient coordinates.
 74    """
 75    import pydicom
 76
 77    seg = pydicom.dcmread(seg_path)
 78    frames = seg.pixel_array
 79    if frames.ndim == 2:  # A segmentation with a single frame.
 80        frames = frames[None]
 81
 82    shared_group = seg.SharedFunctionalGroupsSequence[0]
 83    orientation = np.array([float(v) for v in shared_group.PlaneOrientationSequence[0].ImageOrientationPatient])
 84    row_dir, col_dir = orientation[:3], orientation[3:]
 85    normal = np.cross(row_dir, col_dir)
 86    pixel_measures = shared_group.PixelMeasuresSequence[0]
 87    spacing = [float(v) for v in pixel_measures.PixelSpacing]  # The spacing between rows and between columns.
 88
 89    # The frames may be stored in arbitrary order and frames without any foreground may be skipped,
 90    # so the position of each frame along the slice normal is derived from its patient position.
 91    frame_groups = seg.PerFrameFunctionalGroupsSequence
 92    positions = np.array([[float(v) for v in g.PlanePositionSequence[0].ImagePositionPatient] for g in frame_groups])
 93    projections = positions @ normal
 94    if "SpacingBetweenSlices" in pixel_measures:
 95        slice_spacing = float(pixel_measures.SpacingBetweenSlices)
 96    elif len(projections) > 1:
 97        slice_spacing = np.min(np.diff(np.unique(np.round(projections, 3))))
 98    else:
 99        slice_spacing = float(pixel_measures.SliceThickness)
100    slice_ids = np.round((projections - projections.min()) / slice_spacing).astype("int")
101
102    labels = np.zeros((slice_ids.max() + 1, seg.Rows, seg.Columns), dtype="uint8")
103    for frame, frame_group, slice_id in zip(frames, frame_groups, slice_ids):
104        segment_number = int(frame_group.SegmentIdentificationSequence[0].ReferencedSegmentNumber)
105        labels[slice_id][frame.astype("bool")] = segment_number
106
107    affine = np.eye(4)
108    affine[:3, 0] = normal * slice_spacing
109    affine[:3, 1] = col_dir * spacing[0]
110    affine[:3, 2] = row_dir * spacing[1]
111    affine[:3, 3] = positions[np.argmin(projections)]
112    return labels, affine
113
114
115def _resample_labels(labels, affine, target_shape, target_affine):
116    """Resample a label volume onto the voxel grid of a reference image with nearest neighbor interpolation.
117
118    This is exact if both volumes are stored on the same grid (e.g. a DICOM-SEG object stored on a cropped
119    grid of the reference image) and downsamples segmentations that are stored on a finer grid.
120    """
121    to_label_index = np.linalg.inv(affine) @ target_affine
122    resampled = np.zeros(target_shape, dtype=labels.dtype)
123    yy, xx = np.meshgrid(np.arange(target_shape[1]), np.arange(target_shape[2]), indexing="ij")
124    for z in range(target_shape[0]):
125        target_indices = np.stack([np.full(yy.size, z), yy.ravel(), xx.ravel(), np.ones(yy.size)])
126        indices = np.round(to_label_index[:3] @ target_indices).astype("int")
127        valid = np.all((indices >= 0) & (indices < np.array(labels.shape)[:, None]), axis=0)
128        resampled[z].flat[valid] = labels[tuple(indices[:, valid])]
129    return resampled
130
131
132def _preprocess_adrenal_acc(dicom_dir, csv_path, preprocessed_dir):
133    import h5py
134    import pydicom
135
136    with open(csv_path, "r") as f:
137        seg_series = {row["Subject ID"]: row["Series UID"] for row in csv.DictReader(f) if row["Modality"] == "SEG"}
138
139    os.makedirs(preprocessed_dir, exist_ok=True)
140    for subject_id, seg_uid in tqdm(sorted(seg_series.items()), desc="Preprocess Adrenal-ACC-Ki67-Seg"):
141        out_path = os.path.join(preprocessed_dir, f"{subject_id}.h5")
142        if os.path.exists(out_path):
143            continue
144
145        # The segmentation references the CT series it was created for.
146        seg_path = glob(os.path.join(dicom_dir, seg_uid, "*.dcm"))[0]
147        ct_uid = pydicom.dcmread(seg_path, stop_before_pixels=True).ReferencedSeriesSequence[0].SeriesInstanceUID
148        volume, affine = _load_dicom_volume(os.path.join(dicom_dir, ct_uid))
149        seg_labels, seg_affine = _load_dicom_seg(seg_path)
150        assert seg_labels.max() == 1, f"Expected a single segment in {seg_path}."
151        labels = _resample_labels(seg_labels, seg_affine, volume.shape, affine) * LABEL_IDS["tumor"]
152
153        with h5py.File(out_path, "w") as f:
154            f.create_dataset("raw", data=volume, compression="gzip")
155            f.create_dataset("labels", data=labels, compression="gzip")
156
157
158def get_adrenal_acc_data(path: Union[os.PathLike, str], download: bool = False) -> str:
159    """Download the Adrenal-ACC-Ki67-Seg dataset.
160
161    Args:
162        path: Filepath to a folder where the data is downloaded for further processing.
163        download: Whether to download the data if it is not present.
164
165    Returns:
166        Filepath where the preprocessed data is stored.
167    """
168    preprocessed_dir = os.path.join(path, "preprocessed")
169    if os.path.exists(preprocessed_dir):
170        return preprocessed_dir
171
172    os.makedirs(path, exist_ok=True)
173
174    # Download the DICOM series (CT and SEG) from the TCIA manifest.
175    dicom_dir = os.path.join(path, "dicom")
176    csv_path = os.path.join(path, "adrenal_acc_series")
177    util.download_source_tcia(
178        path=os.path.join(path, "Adrenal-ACC-Ki67-Seg_v1.tcia"), url=URL, dst=dicom_dir,
179        csv_filename=csv_path, download=download,
180    )
181
182    _preprocess_adrenal_acc(dicom_dir, f"{csv_path}.csv", preprocessed_dir)
183    return preprocessed_dir
184
185
186def get_adrenal_acc_paths(path: Union[os.PathLike, str], download: bool = False) -> List[str]:
187    """Get paths to the Adrenal-ACC-Ki67-Seg data.
188
189    Args:
190        path: Filepath to a folder where the data is downloaded for further processing.
191        download: Whether to download the data if it is not present.
192
193    Returns:
194        List of filepaths for the hdf5 files, which contain the image data ('raw') and the label data ('labels').
195    """
196    data_dir = get_adrenal_acc_data(path, download)
197    volume_paths = natsorted(glob(os.path.join(data_dir, "*.h5")))
198    return volume_paths
199
200
201def get_adrenal_acc_dataset(
202    path: Union[os.PathLike, str],
203    patch_shape: Tuple[int, ...],
204    resize_inputs: bool = False,
205    download: bool = False,
206    **kwargs
207) -> Dataset:
208    """Get the Adrenal-ACC-Ki67-Seg dataset for adrenal tumor segmentation.
209
210    Args:
211        path: Filepath to a folder where the data is downloaded for further processing.
212        patch_shape: The patch shape to use for training.
213        resize_inputs: Whether to resize inputs to the desired patch shape.
214        download: Whether to download the data if it is not present.
215        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
216
217    Returns:
218        The segmentation dataset.
219    """
220    volume_paths = get_adrenal_acc_paths(path, download)
221
222    if resize_inputs:
223        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
224        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
225            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
226        )
227
228    return torch_em.default_segmentation_dataset(
229        raw_paths=volume_paths,
230        raw_key="raw",
231        label_paths=volume_paths,
232        label_key="labels",
233        patch_shape=patch_shape,
234        is_seg_dataset=True,
235        **kwargs
236    )
237
238
239def get_adrenal_acc_loader(
240    path: Union[os.PathLike, str],
241    batch_size: int,
242    patch_shape: Tuple[int, ...],
243    resize_inputs: bool = False,
244    download: bool = False,
245    **kwargs
246) -> DataLoader:
247    """Get the Adrenal-ACC-Ki67-Seg dataloader for adrenal tumor segmentation.
248
249    Args:
250        path: Filepath to a folder where the data is downloaded for further processing.
251        batch_size: The batch size for training.
252        patch_shape: The patch shape to use for training.
253        resize_inputs: Whether to resize inputs to the desired patch shape.
254        download: Whether to download the data if it is not present.
255        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
256
257    Returns:
258        The DataLoader.
259    """
260    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
261    dataset = get_adrenal_acc_dataset(path, patch_shape, resize_inputs, download, **ds_kwargs)
262    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)