torch_em.data.datasets.histopathology.bc_multiset

  1"""The BC-MultiSet dataset contains nucleus annotations for H&E stained images of invasive breast carcinoma.
  2
  3The dataset holds 5,586 patches of 256 x 256 pixels at 40x magnification from 50 patients, split into
  4four cross-validation folds. Two board-certified pathologists labeled the nuclei of three cell types (tumor,
  5fibroblast and inflammatory) and reconciled them by consensus. It also provides patient-level clinical and
  6molecular data (ER, PR, HER2, Ki67, TILs and more), which this module does not use.
  7NOTE: The patches are cut from larger crops with a stride of 192 pixels, so neighboring patches overlap.
  8NOTE: Not every nucleus is annotated: some nuclei in the images carry no label and become background.
  9NOTE: The label files are pickled numpy dicts. This module reads them with a restricted unpickler that only
 10accepts numpy arrays, so that no code from the files is executed.
 11
 12The dataset is located at https://doi.org/10.5281/zenodo.20717567 (CC BY 4.0).
 13Please cite it if you use this dataset for your research.
 14"""
 15
 16import io
 17import os
 18import pickle
 19import zipfile
 20from glob import glob
 21from tqdm import tqdm
 22from natsort import natsorted
 23from typing import Union, Tuple, List, Literal, Optional, Sequence
 24
 25import numpy as np
 26import imageio.v3 as imageio
 27
 28from torch.utils.data import Dataset, DataLoader
 29
 30import torch_em
 31
 32from .. import util
 33
 34
 35URL = "https://zenodo.org/api/records/20717567/files/BC-MultiSet-50.zip/content"
 36CHECKSUM = "1da5d93f397c80f0cf7e41adaa5e432d5164a0918588fc2629ee9f5c1fdbe208"
 37FOLDS = (1, 2, 3, 4)
 38
 39# The only globals that the label pickles need.
 40_ALLOWED_GLOBALS = {
 41    ("numpy.core.multiarray", "_reconstruct"), ("numpy._core.multiarray", "_reconstruct"),
 42    ("numpy", "ndarray"), ("numpy", "dtype"),
 43}
 44
 45
 46class _NumpyOnlyUnpickler(pickle.Unpickler):
 47    def find_class(self, module, name):
 48        if (module, name) not in _ALLOWED_GLOBALS:
 49            raise pickle.UnpicklingError(f"The global '{module}.{name}' is not allowed in a label file.")
 50        if name == "_reconstruct":
 51            try:
 52                from numpy._core.multiarray import _reconstruct
 53            except ImportError:  # numpy < 2
 54                from numpy.core.multiarray import _reconstruct
 55            return _reconstruct
 56        return getattr(np, name)
 57
 58
 59def _read_labels(data):
 60    import numpy.lib.format as npy_format
 61
 62    stream = io.BytesIO(data)
 63    version = npy_format.read_magic(stream)
 64    read_header = npy_format.read_array_header_1_0 if version == (1, 0) else npy_format.read_array_header_2_0
 65    shape, _, dtype = read_header(stream)
 66    assert shape == () and dtype == np.dtype("O"), (shape, dtype)
 67    # The pickle holds a 0-d object array, which wraps the dict of label maps.
 68    labels = _NumpyOnlyUnpickler(stream).load().item()
 69    return labels["inst_map"], labels["type_map"]
 70
 71
 72def _create_h5_files(path, zip_path):
 73    import h5py
 74
 75    with zipfile.ZipFile(zip_path) as archive:
 76        for fold in FOLDS:
 77            out_dir = os.path.join(path, "preprocessed", f"fold{fold}")
 78            os.makedirs(out_dir, exist_ok=True)
 79            prefix = f"BC-MultiSet-50/fold{fold}/labels/"
 80            names = [n for n in archive.namelist() if n.startswith(prefix) and n.endswith(".npy")]
 81            for name in tqdm(names, desc=f"Preprocess BC-MultiSet fold {fold}"):
 82                stem = os.path.splitext(os.path.basename(name))[0]
 83                image_name = f"BC-MultiSet-50/fold{fold}/images/{stem}.png"
 84                if image_name not in archive.NameToInfo:
 85                    continue
 86                image = imageio.imread(archive.read(image_name), extension=".png")[..., :3]
 87                instances, types = _read_labels(archive.read(name))
 88                assert image.shape[:2] == instances.shape == types.shape, stem
 89
 90                with h5py.File(os.path.join(out_dir, f"{stem}.h5"), "w") as f:
 91                    f.create_dataset("raw", data=image.transpose(2, 0, 1), compression="gzip")
 92                    f.create_dataset("labels/instances", data=instances.astype("uint32"), compression="gzip")
 93                    f.create_dataset("labels/semantic", data=types.astype("uint8"), compression="gzip")
 94
 95
 96def get_bc_multiset_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 97    """Download the BC-MultiSet dataset.
 98
 99    Args:
100        path: Filepath to a folder where the downloaded data is saved.
101        download: Whether to download the data if it is not present.
102
103    Returns:
104        Filepath where the preprocessed data is stored.
105    """
106    data_dir = os.path.join(path, "preprocessed")
107    if os.path.exists(data_dir):
108        return data_dir
109
110    os.makedirs(path, exist_ok=True)
111    zip_path = os.path.join(path, "BC-MultiSet-50.zip")
112    util.download_source(zip_path, URL, download, checksum=CHECKSUM)
113    _create_h5_files(path, zip_path)
114    return data_dir
115
116
117def get_bc_multiset_paths(
118    path: Union[os.PathLike, str], folds: Optional[Sequence[int]] = None, download: bool = False
119) -> List[str]:
120    """Get paths to the BC-MultiSet data.
121
122    Args:
123        path: Filepath to a folder where the downloaded data is saved.
124        folds: The cross-validation folds, out of 1, 2, 3 and 4. By default, all folds are returned.
125        download: Whether to download the data if it is not present.
126
127    Returns:
128        List of filepaths for the input data.
129    """
130    data_dir = get_bc_multiset_data(path, download)
131    folds = FOLDS if folds is None else folds
132    assert all(fold in FOLDS for fold in folds), folds
133    return [p for fold in folds for p in natsorted(glob(os.path.join(data_dir, f"fold{fold}", "*.h5")))]
134
135
136def get_bc_multiset_dataset(
137    path: Union[os.PathLike, str],
138    patch_shape: Tuple[int, int],
139    folds: Optional[Sequence[int]] = None,
140    label_choice: Literal["instances", "semantic"] = "instances",
141    download: bool = False,
142    **kwargs
143) -> Dataset:
144    """Get the BC-MultiSet dataset for nucleus segmentation.
145
146    Args:
147        path: Filepath to a folder where the downloaded data is saved.
148        patch_shape: The patch shape to use for training.
149        folds: The cross-validation folds, out of 1, 2, 3 and 4. By default, all folds are used.
150        label_choice: The choice of labels, 'instances' or 'semantic' (the three cell types).
151        download: Whether to download the data if it is not present.
152        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
153
154    Returns:
155        The segmentation dataset.
156    """
157    assert label_choice in ("instances", "semantic"), label_choice
158    paths = get_bc_multiset_paths(path, folds, download)
159    return torch_em.default_segmentation_dataset(
160        raw_paths=paths,
161        raw_key="raw",
162        label_paths=paths,
163        label_key=f"labels/{label_choice}",
164        patch_shape=patch_shape,
165        ndim=2,
166        with_channels=True,
167        **kwargs
168    )
169
170
171def get_bc_multiset_loader(
172    path: Union[os.PathLike, str],
173    batch_size: int,
174    patch_shape: Tuple[int, int],
175    folds: Optional[Sequence[int]] = None,
176    label_choice: Literal["instances", "semantic"] = "instances",
177    download: bool = False,
178    **kwargs
179) -> DataLoader:
180    """Get the BC-MultiSet dataloader for nucleus segmentation.
181
182    Args:
183        path: Filepath to a folder where the downloaded data is saved.
184        batch_size: The batch size for training.
185        patch_shape: The patch shape to use for training.
186        folds: The cross-validation folds, out of 1, 2, 3 and 4. By default, all folds are used.
187        label_choice: The choice of labels, 'instances' or 'semantic' (the three cell types).
188        download: Whether to download the data if it is not present.
189        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
190
191    Returns:
192        The DataLoader.
193    """
194    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
195    dataset = get_bc_multiset_dataset(path, patch_shape, folds, label_choice, download, **ds_kwargs)
196    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)