torch_em.data.datasets.medical.promise12

  1"""The PROMISE12 dataset contains annotations for prostate segmentation in T2-weighted MRI.
  2
  3The data comes from the MICCAI 2012 'Prostate MR Image Segmentation' challenge (https://promise12.grand-challenge.org/)
  4and covers 50 labeled training cases from four centers with different scanners, field strengths and protocols
  5(with and without an endorectal coil). The label ids are: background: 0 and prostate: 1.
  6The challenge test cases are not labeled and are therefore not exposed here.
  7
  8The original MetaImage ('.mhd' / '.raw') volumes are converted to hdf5 once and then loaded from there.
  9
 10The dataset is located at https://zenodo.org/records/8026660.
 11
 12This dataset is from the publication https://doi.org/10.1016/j.media.2013.12.002.
 13Please cite it if you use this dataset in your research.
 14"""
 15
 16import os
 17from glob import glob
 18from tqdm import tqdm
 19from natsort import natsorted
 20from typing import Union, Tuple, List
 21
 22import numpy as np
 23
 24from torch.utils.data import Dataset, DataLoader
 25
 26import torch_em
 27
 28from .. import util
 29
 30
 31URL = "https://zenodo.org/records/8026660/files/training_data.zip?download=1"
 32CHECKSUM = "150287d0c74cd0105d8b70b43af4a7bf4f1fd3e1748829779177a0ccaf948f45"
 33
 34METAIMAGE_DTYPES = {
 35    "MET_CHAR": "int8", "MET_UCHAR": "uint8", "MET_SHORT": "int16", "MET_USHORT": "uint16",
 36    "MET_INT": "int32", "MET_UINT": "uint32", "MET_FLOAT": "float32", "MET_DOUBLE": "float64",
 37}
 38
 39
 40def _read_metaimage(path):
 41    """Read an uncompressed MetaImage ('.mhd' + '.raw') volume with numpy, returning it in 'zyx' order."""
 42    header = {}
 43    with open(path, "r") as f:
 44        for line in f:
 45            if "=" in line:
 46                key, value = line.split("=", 1)
 47                header[key.strip()] = value.strip()
 48
 49    if header.get("CompressedData", "False") == "True":
 50        raise NotImplementedError(f"Compressed MetaImage data is not supported: '{path}'.")
 51
 52    shape = tuple(int(s) for s in header["DimSize"].split())[::-1]
 53    dtype = np.dtype(METAIMAGE_DTYPES[header["ElementType"]])
 54    if header.get("BinaryDataByteOrderMSB", "False") == "True":
 55        dtype = dtype.newbyteorder(">")
 56
 57    raw_path = os.path.join(os.path.dirname(path), header["ElementDataFile"])
 58    return np.fromfile(raw_path, dtype=dtype).reshape(shape)
 59
 60
 61def get_promise12_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 62    """Download the PROMISE12 dataset and convert it to hdf5.
 63
 64    Args:
 65        path: Filepath to a folder where the data is downloaded for further processing.
 66        download: Whether to download the data if it is not present.
 67
 68    Returns:
 69        Filepath to the folder with the converted hdf5 volumes.
 70    """
 71    data_dir = os.path.join(path, "preprocessed")
 72    if os.path.exists(data_dir):
 73        return data_dir
 74
 75    os.makedirs(path, exist_ok=True)
 76
 77    zip_path = os.path.join(path, "training_data.zip")
 78    util.download_source(path=zip_path, url=URL, download=download, checksum=CHECKSUM)
 79    raw_dir = os.path.join(path, "training_data")
 80    util.unzip(zip_path=zip_path, dst=raw_dir)
 81
 82    import h5py
 83
 84    os.makedirs(data_dir, exist_ok=True)
 85    label_paths = natsorted(glob(os.path.join(raw_dir, "Case*_segmentation.mhd")))
 86    for label_path in tqdm(label_paths, desc="Converting PROMISE12 volumes to hdf5"):
 87        raw_path = label_path.replace("_segmentation.mhd", ".mhd")
 88        raw = _read_metaimage(raw_path)
 89        labels = _read_metaimage(label_path).astype("uint8")
 90        assert raw.shape == labels.shape, f"Shape mismatch for '{raw_path}': {raw.shape} vs. {labels.shape}."
 91
 92        case_id = os.path.basename(raw_path).replace(".mhd", "")
 93        with h5py.File(os.path.join(data_dir, f"{case_id}.h5"), "w") as f:
 94            f.create_dataset("raw", data=raw, compression="gzip")
 95            f.create_dataset("labels", data=labels, compression="gzip")
 96
 97    return data_dir
 98
 99
100def get_promise12_paths(path: Union[os.PathLike, str], download: bool = False) -> List[str]:
101    """Get paths to the PROMISE12 data.
102
103    Args:
104        path: Filepath to a folder where the data is downloaded for further processing.
105        download: Whether to download the data if it is not present.
106
107    Returns:
108        List of filepaths for the hdf5 volumes, which store the image data at 'raw' and the label data at 'labels'.
109    """
110    data_dir = get_promise12_data(path, download)
111    volume_paths = natsorted(glob(os.path.join(data_dir, "Case*.h5")))
112    assert len(volume_paths) > 0, f"No PROMISE12 volumes found at '{data_dir}'."
113    return volume_paths
114
115
116def get_promise12_dataset(
117    path: Union[os.PathLike, str],
118    patch_shape: Tuple[int, ...],
119    resize_inputs: bool = False,
120    download: bool = False,
121    **kwargs
122) -> Dataset:
123    """Get the PROMISE12 dataset for prostate segmentation.
124
125    Args:
126        path: Filepath to a folder where the data is downloaded for further processing.
127        patch_shape: The patch shape to use for training.
128        resize_inputs: Whether to resize inputs to the desired patch shape.
129        download: Whether to download the data if it is not present.
130        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
131
132    Returns:
133        The segmentation dataset.
134    """
135    volume_paths = get_promise12_paths(path, download)
136
137    if resize_inputs:
138        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
139        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
140            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
141        )
142
143    return torch_em.default_segmentation_dataset(
144        raw_paths=volume_paths,
145        raw_key="raw",
146        label_paths=volume_paths,
147        label_key="labels",
148        patch_shape=patch_shape,
149        is_seg_dataset=True,
150        **kwargs
151    )
152
153
154def get_promise12_loader(
155    path: Union[os.PathLike, str],
156    batch_size: int,
157    patch_shape: Tuple[int, ...],
158    resize_inputs: bool = False,
159    download: bool = False,
160    **kwargs
161) -> DataLoader:
162    """Get the PROMISE12 dataloader for prostate segmentation.
163
164    Args:
165        path: Filepath to a folder where the data is downloaded for further processing.
166        batch_size: The batch size for training.
167        patch_shape: The patch shape to use for training.
168        resize_inputs: Whether to resize inputs to the desired patch shape.
169        download: Whether to download the data if it is not present.
170        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
171
172    Returns:
173        The DataLoader.
174    """
175    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
176    dataset = get_promise12_dataset(path, patch_shape, resize_inputs, download, **ds_kwargs)
177    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)