torch_em.data.datasets.medical.advanced_mri_breast_lesions

  1"""The Advanced-MRI-Breast-Lesions dataset contains annotations for breast lesions in MRI.
  2
  3The dataset consists of dynamic contrast-enhanced breast MRI series, each with an expert DICOM-SEG
  4object marking the lesion mass or masses (multiple, distinctly labelled instances for a multifocal
  5case), paired here with the exact series it references.
  6
  7NOTE: This requires the pydicom python package.
  8
  9The dataset is located at https://doi.org/10.7937/C7X1-YN57 and is distributed under the
 10CC BY 4.0 license.
 11Please cite it if you use this dataset in your research.
 12"""
 13
 14import os
 15import json
 16from glob import glob
 17from tqdm import tqdm
 18from natsort import natsorted
 19from typing import Union, Tuple, List
 20
 21from torch.utils.data import Dataset, DataLoader
 22
 23import torch_em
 24
 25from .adrenal_acc import _load_dicom_volume, _load_dicom_seg, _resample_labels
 26from .. import util
 27
 28
 29COLLECTION = "Advanced-MRI-Breast-Lesions"
 30
 31
 32def _get_series_metadata(path, download):
 33    """Get the metadata of all series in the collection from the NBIA REST API."""
 34    import requests
 35
 36    metadata_path = os.path.join(path, "advanced_mri_breast_lesions_series.json")
 37    if not os.path.exists(metadata_path):
 38        if not download:
 39            raise RuntimeError(f"Cannot find the data at {path}, but download was set to False.")
 40        response = requests.get(f"{util.NBIA_API_URL}getSeries", params={"Collection": COLLECTION})
 41        response.raise_for_status()
 42        with open(metadata_path, "w") as f:
 43            json.dump(response.json(), f, indent=2)
 44
 45    with open(metadata_path, "r") as f:
 46        return json.load(f)
 47
 48
 49def _referenced_series_uid(seg):
 50    return str(seg.ReferencedSeriesSequence[0].SeriesInstanceUID)
 51
 52
 53def _preprocess_advanced_mri_breast_lesions(dicom_dir, series_metadata, preprocessed_dir):
 54    import h5py
 55    import pydicom
 56
 57    seg_series = [series for series in series_metadata if series.get("Modality") == "SEG"]
 58
 59    os.makedirs(preprocessed_dir, exist_ok=True)
 60    for series in tqdm(seg_series, desc="Preprocess Advanced-MRI-Breast-Lesions"):
 61        seg_paths = glob(os.path.join(dicom_dir, series["SeriesInstanceUID"], "*.dcm"))
 62        if not seg_paths:
 63            continue
 64
 65        out_path = os.path.join(preprocessed_dir, f"{series['SeriesInstanceUID']}.h5")
 66        if os.path.exists(out_path):
 67            continue
 68
 69        seg = pydicom.dcmread(seg_paths[0])
 70        mr_dir = os.path.join(dicom_dir, _referenced_series_uid(seg))
 71        if not glob(os.path.join(mr_dir, "*.dcm")):
 72            continue
 73
 74        volume, mr_affine = _load_dicom_volume(mr_dir)
 75        seg_labels, seg_affine = _load_dicom_seg(seg_paths[0])
 76        labels = _resample_labels(seg_labels, seg_affine, volume.shape, mr_affine)
 77        if labels.max() == 0:
 78            continue
 79
 80        with h5py.File(out_path, "w") as f:
 81            f.create_dataset("raw", data=volume, compression="gzip")
 82            f.create_dataset("labels", data=labels, compression="gzip")
 83
 84
 85def get_advanced_mri_breast_lesions_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 86    """Download the Advanced-MRI-Breast-Lesions dataset.
 87
 88    Args:
 89        path: Filepath to a folder where the data is downloaded for further processing.
 90        download: Whether to download the data if it is not present.
 91
 92    Returns:
 93        Filepath where the preprocessed data is stored.
 94    """
 95    # NOTE: The preprocessing below skips volumes that were converted already, so an interrupted run resumes.
 96    preprocessed_dir = os.path.join(path, "preprocessed")
 97
 98    os.makedirs(path, exist_ok=True)
 99    series_metadata = _get_series_metadata(path, download)
100
101    seg_uids = [series["SeriesInstanceUID"] for series in series_metadata if series.get("Modality") == "SEG"]
102
103    dicom_dir = os.path.join(path, "dicom")
104    if download:  # The SEG series are downloaded first, so the MR series they reference can be found.
105        util.download_tcia_series(
106            seg_uids, dst=dicom_dir, csv_filename=os.path.join(path, "advanced_mri_breast_lesions_seg")
107        )
108
109        import pydicom
110        mr_uids = set()
111        for uid in seg_uids:
112            seg_paths = glob(os.path.join(dicom_dir, uid, "*.dcm"))
113            if seg_paths:
114                seg = pydicom.dcmread(seg_paths[0])
115                mr_uids.add(_referenced_series_uid(seg))
116        util.download_tcia_series(
117            sorted(mr_uids), dst=dicom_dir, csv_filename=os.path.join(path, "advanced_mri_breast_lesions_mr")
118        )
119    elif not glob(os.path.join(dicom_dir, "*", "*.dcm")):
120        raise RuntimeError(f"Cannot find the data at {path}, but download was set to False.")
121
122    _preprocess_advanced_mri_breast_lesions(dicom_dir, series_metadata, preprocessed_dir)
123    return preprocessed_dir
124
125
126def get_advanced_mri_breast_lesions_paths(path: Union[os.PathLike, str], download: bool = False) -> List[str]:
127    """Get paths to the Advanced-MRI-Breast-Lesions data.
128
129    Args:
130        path: Filepath to a folder where the data is downloaded for further processing.
131        download: Whether to download the data if it is not present.
132
133    Returns:
134        List of filepaths for the stored data.
135    """
136    preprocessed_dir = get_advanced_mri_breast_lesions_data(path, download)
137    volume_paths = natsorted(glob(os.path.join(preprocessed_dir, "*.h5")))
138    assert len(volume_paths) > 0, f"Could not find any preprocessed volumes in '{preprocessed_dir}'."
139    return volume_paths
140
141
142def get_advanced_mri_breast_lesions_dataset(
143    path: Union[os.PathLike, str],
144    patch_shape: Tuple[int, ...],
145    resize_inputs: bool = False,
146    download: bool = False,
147    **kwargs
148) -> Dataset:
149    """Get the Advanced-MRI-Breast-Lesions dataset for breast lesion segmentation.
150
151    Args:
152        path: Filepath to a folder where the data is downloaded for further processing.
153        patch_shape: The patch shape to use for training.
154        resize_inputs: Whether to resize inputs to the desired patch shape.
155        download: Whether to download the data if it is not present.
156        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
157
158    Returns:
159        The segmentation dataset.
160    """
161    volume_paths = get_advanced_mri_breast_lesions_paths(path, download)
162
163    if resize_inputs:
164        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
165        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
166            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
167        )
168
169    return torch_em.default_segmentation_dataset(
170        raw_paths=volume_paths,
171        raw_key="raw",
172        label_paths=volume_paths,
173        label_key="labels",
174        patch_shape=patch_shape,
175        is_seg_dataset=True,
176        **kwargs
177    )
178
179
180def get_advanced_mri_breast_lesions_loader(
181    path: Union[os.PathLike, str],
182    batch_size: int,
183    patch_shape: Tuple[int, ...],
184    resize_inputs: bool = False,
185    download: bool = False,
186    **kwargs
187) -> DataLoader:
188    """Get the Advanced-MRI-Breast-Lesions dataloader for breast lesion segmentation.
189
190    Args:
191        path: Filepath to a folder where the data is downloaded for further processing.
192        batch_size: The batch size for training.
193        patch_shape: The patch shape to use for training.
194        resize_inputs: Whether to resize inputs to the desired patch shape.
195        download: Whether to download the data if it is not present.
196        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
197
198    Returns:
199        The DataLoader.
200    """
201    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
202    dataset = get_advanced_mri_breast_lesions_dataset(path, patch_shape, resize_inputs, download, **ds_kwargs)
203    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)