torch_em.data.datasets.histopathology.deepliif

  1"""DeepLIIF contains annotations for nucleus segmentation and classification in IHC images
  2of lung, bladder and breast cancer tissue.
  3
  4The lung and bladder tissue shares one archive per split. The breast cancer tissue, which is
  5stained for Ki67, stems from an earlier release of the data and has its own archives. It only
  6covers the train and the val split.
  7
  8Every image comes as six 512x512 panels of the same region, which are stitched next to each other:
  9the IHC input, the hematoxylin channel and the multiplex immunofluorescence modalities
 10(mpIF DAPI, mpIF Lap2 and the mpIF marker), followed by the segmentation mask.
 11The mask marks cells with positive protein expression in red and negative cells in blue.
 12A green boundary surrounds each cell, so that this module can separate touching cells
 13and derive instance labels from the mask.
 14NOTE: The labels keep every cell of the mask. A few cells are only a few pixels large,
 15since saving the masks as png introduced compression artifacts.
 16NOTE: The segmentations align with the mpIF DAPI channel. They are slightly misaligned with
 17all other modalities.
 18
 19The data is hosted at https://zenodo.org/records/4751737 and licensed under CC-BY-4.0.
 20This dataset is from the publication https://doi.org/10.1038/s42256-022-00471-x.
 21Please cite it if you use this dataset for your research.
 22"""
 23
 24import os
 25from glob import glob
 26from natsort import natsorted
 27from typing import Union, Tuple, List, Literal, Optional
 28
 29import numpy as np
 30import imageio.v3 as imageio
 31from skimage.measure import label as connected_components
 32from skimage.segmentation import watershed, find_boundaries, relabel_sequential
 33
 34from torch.utils.data import Dataset, DataLoader
 35
 36import torch_em
 37
 38from .. import util
 39
 40
 41URL = "https://zenodo.org/records/4751737/files/{filename}?download=1"
 42
 43# The lung and bladder tissue shares one archive per split, the breast cancer tissue has its own.
 44FILENAMES = {
 45    ("lung_bladder", "train"): "DeepLIIF_Training_Set.zip",
 46    ("lung_bladder", "val"): "DeepLIIF_Validation_Set.zip",
 47    ("lung_bladder", "test"): "DeepLIIF_Testing_Set.zip",
 48    ("breast", "train"): "BC-DeepLIIF_Training_Set.zip",
 49    ("breast", "val"): "BC-DeepLIIF_Validation_Set.zip",
 50}
 51
 52CHECKSUMS = {
 53    ("lung_bladder", "train"): "704b75ad9d15d5c5ffda8c48a53a7037544208bbcb7a1f3bbb2b1ce9c4f63d02",
 54    ("lung_bladder", "val"): "e8dcec56d4cb44a7037170060e39ec503ba3b70f817abd00cb36fb86eb4e7e1a",
 55    ("lung_bladder", "test"): "7663611b1274049b677dd18deee8fbd1e44370884ec1f4d8d7196924d2a7acec",
 56    ("breast", "train"): "598bef02f1dcc54f888976ad5ed267e1672816e9326c2011a35c68d5b0167d8e",
 57    ("breast", "val"): "349278854a95b0a31790e81acd21f87c7afaa8f3641ce5fb4c47b09987e45b23",
 58}
 59
 60SPLITS = ["train", "val", "test"]
 61
 62TISSUES = ["lung", "bladder", "breast"]
 63
 64# The lung and bladder images carry their tissue in the filename, the breast cancer images do not.
 65FILE_PREFIXES = {"lung": "Lung_", "bladder": "Bladder_"}
 66
 67# The image panels, in the order in which they are stitched together. The mask follows them.
 68MODALITIES = ["ihc", "hematoxylin", "dapi", "lap2", "marker"]
 69
 70# Half of the boundary that the mask draws around each cell, measured per archive. The watershed
 71# grows the cells over the full boundary, so shrinking them by this puts their edge in its middle.
 72EROSION = {"lung_bladder": 2, "breast": 3}
 73
 74
 75def _get_archive(tissue, split):
 76    """Map a tissue to the archive that holds it for a split."""
 77    key = ("breast" if tissue == "breast" else "lung_bladder", split)
 78    if key not in FILENAMES:
 79        raise ValueError(f"The '{tissue}' tissue has no '{split}' split.")
 80    return key
 81
 82
 83def _get_labels(mask, n_erode):
 84    """Derive the instance and the semantic labels from the segmentation panel.
 85
 86    The panel paints cells with positive expression red and negative cells blue, and draws a
 87    green boundary around every cell. The cell interiors are separate already, so they seed a
 88    watershed that grows them back over the boundary. This recovers the full extent of a cell,
 89    which taking the interiors alone would shrink by the width of the boundary. Eroding the
 90    result afterwards moves the edge of a cell from the outside of the boundary to its middle.
 91
 92    The class of a cell follows from the color of its interior, so that the semantic labels
 93    and the instances agree on where a cell ends.
 94    """
 95    positive, negative = mask[..., 0] > 127, mask[..., 2] > 127
 96    interior = positive | negative
 97    foreground = interior | (mask[..., 1] > 127)
 98
 99    seeds = connected_components(interior)
100    instances = watershed(mask[..., 1], markers=seeds, mask=foreground)
101
102    # Erode every cell on its own, so that cells which touch stay apart.
103    for _ in range(n_erode):
104        instances[find_boundaries(instances, mode="inner")] = 0
105
106    # The erosion removes the smallest cells, so make the labels consecutive again.
107    instances = relabel_sequential(instances)[0].astype("uint16")
108
109    # A cell is positive if its interior holds more red than blue pixels.
110    n_labels = int(instances.max()) + 1
111    n_positive = np.bincount(instances[positive], minlength=n_labels)
112    n_negative = np.bincount(instances[negative], minlength=n_labels)
113
114    classes = np.where(n_positive >= n_negative, 2, 1).astype("uint8")
115    classes[(n_positive + n_negative) == 0] = 0
116    classes[0] = 0
117    semantic = classes[instances]
118
119    return instances, semantic
120
121
122def _preprocess_data(input_dir, data_dir, n_erode):
123    import h5py
124
125    os.makedirs(data_dir, exist_ok=True)
126    image_paths = natsorted(glob(os.path.join(input_dir, "*.png")))
127    if not image_paths:
128        raise RuntimeError(f"Could not find the images in {input_dir}.")
129
130    for image_path in image_paths:
131        fname = os.path.splitext(os.path.basename(image_path))[0]
132        out_path = os.path.join(data_dir, f"{fname}.h5")
133        if os.path.exists(out_path):
134            continue
135
136        panels = np.split(imageio.imread(image_path)[..., :3], len(MODALITIES) + 1, axis=1)
137        instances, semantic = _get_labels(panels[-1], n_erode)
138
139        with h5py.File(out_path, "a") as f:
140            for modality, panel in zip(MODALITIES, panels):
141                f.create_dataset(f"raw/{modality}", data=panel.transpose(2, 0, 1), compression="gzip")
142
143            f.create_dataset("labels/instances", data=instances, compression="gzip")
144            f.create_dataset("labels/semantic", data=semantic, compression="gzip")
145
146
147def _get_tissues(tissue, split):
148    """Resolve the tissue argument to the tissues that the split actually holds."""
149    if tissue is None:
150        return [name for name in TISSUES if ("breast" if name == "breast" else "lung_bladder", split) in FILENAMES]
151
152    tissues = [tissue] if isinstance(tissue, str) else list(tissue)
153    for name in tissues:
154        if name not in TISSUES:
155            raise ValueError(f"'{name}' is not a valid tissue. Choose one of {TISSUES}.")
156    return tissues
157
158
159def get_deepliif_data(
160    path: Union[os.PathLike, str],
161    split: Literal["train", "val", "test"],
162    tissue: Optional[Union[str, List[str]]] = None,
163    download: bool = False,
164) -> List[str]:
165    """Download the DeepLIIF dataset for one split.
166
167    Args:
168        path: The folder where the function stores the data.
169        split: The split of the dataset. Either 'train', 'val' or 'test'.
170        tissue: The tissue to download. See `TISSUES` for the valid choices. By default this
171            downloads every tissue that the split holds.
172        download: Whether to download the data if it is not present.
173
174    Returns:
175        The list of filepaths to the folders with the prepared data.
176    """
177    if split not in SPLITS:
178        raise ValueError(f"'{split}' is not a valid split. Choose one of {SPLITS}.")
179
180    data_dirs = []
181    for archive in dict.fromkeys(_get_archive(name, split) for name in _get_tissues(tissue, split)):
182        archive_dir = os.path.join(path, "_".join(archive))
183        data_dir = os.path.join(archive_dir, "data")
184        data_dirs.append(data_dir)
185        if glob(os.path.join(data_dir, "*.h5")):
186            continue
187
188        os.makedirs(archive_dir, exist_ok=True)
189        filename = FILENAMES[archive]
190        input_dir = os.path.join(archive_dir, os.path.splitext(filename)[0])
191        if not os.path.exists(input_dir):
192            zip_path = os.path.join(archive_dir, filename)
193            util.download_source(
194                path=zip_path, url=URL.format(filename=filename), download=download, checksum=CHECKSUMS[archive]
195            )
196            util.unzip(zip_path=zip_path, dst=archive_dir)
197
198        _preprocess_data(input_dir, data_dir, EROSION[archive[0]])
199
200    return data_dirs
201
202
203def get_deepliif_paths(
204    path: Union[os.PathLike, str],
205    split: Literal["train", "val", "test"],
206    tissue: Optional[Union[str, List[str]]] = None,
207    download: bool = False,
208) -> List[str]:
209    """Get the paths to the DeepLIIF data.
210
211    Args:
212        path: The folder where the function stores the data.
213        split: The split of the dataset. Either 'train', 'val' or 'test'.
214        tissue: The tissue to use. See `TISSUES` for the valid choices. By default this uses
215            every tissue that the split holds.
216        download: Whether to download the data if it is not present.
217
218    Returns:
219        The list of filepaths to the input data.
220    """
221    tissues = _get_tissues(tissue, split)
222    get_deepliif_data(path, split, tissues, download)
223
224    volume_paths = []
225    for name in tissues:
226        data_dir = os.path.join(path, "_".join(_get_archive(name, split)), "data")
227        prefix = FILE_PREFIXES.get(name, "")
228        volume_paths.extend(natsorted(glob(os.path.join(data_dir, f"{prefix}*.h5"))))
229
230    assert len(volume_paths) > 0, f"Could not find data for the split '{split}' and the tissue '{tissue}'."
231    return natsorted(volume_paths)
232
233
234def get_deepliif_dataset(
235    path: Union[os.PathLike, str],
236    patch_shape: Tuple[int, int],
237    split: Literal["train", "val", "test"],
238    tissue: Optional[Union[str, List[str]]] = None,
239    modality: Literal["ihc", "hematoxylin", "dapi", "lap2", "marker"] = "ihc",
240    label_choice: Literal["instances", "semantic"] = "instances",
241    download: bool = False,
242    **kwargs
243) -> Dataset:
244    """Get the DeepLIIF dataset for nucleus segmentation and classification in IHC images.
245
246    Args:
247        path: The folder where the function stores the data.
248        patch_shape: The patch shape to use for training.
249        split: The split of the dataset. Either 'train', 'val' or 'test'.
250        tissue: The tissue to use. See `TISSUES` for the valid choices. By default this uses
251            every tissue that the split holds. Note that 'breast' has no test split.
252        modality: The image modality to use as input. See `MODALITIES` for the valid choices.
253        label_choice: The choice of labels. Either 'instances' for the nucleus instances, or
254            'semantic' for the classification into background, negative cells and positive cells.
255        download: Whether to download the data if it is not present.
256        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
257
258    Returns:
259        The segmentation dataset.
260    """
261    if modality not in MODALITIES:
262        raise ValueError(f"'{modality}' is not a valid modality. Choose one of {MODALITIES}.")
263
264    if label_choice not in ["instances", "semantic"]:
265        raise ValueError(f"'{label_choice}' is not a valid label choice. Choose 'instances' or 'semantic'.")
266
267    volume_paths = get_deepliif_paths(path, split, tissue, download)
268    kwargs = util.update_kwargs(kwargs, "with_channels", True)
269
270    return torch_em.default_segmentation_dataset(
271        raw_paths=volume_paths,
272        raw_key=f"raw/{modality}",
273        label_paths=volume_paths,
274        label_key=f"labels/{label_choice}",
275        patch_shape=patch_shape,
276        is_seg_dataset=True,
277        ndim=2,
278        **kwargs
279    )
280
281
282def get_deepliif_loader(
283    path: Union[os.PathLike, str],
284    batch_size: int,
285    patch_shape: Tuple[int, int],
286    split: Literal["train", "val", "test"],
287    tissue: Optional[Union[str, List[str]]] = None,
288    modality: Literal["ihc", "hematoxylin", "dapi", "lap2", "marker"] = "ihc",
289    label_choice: Literal["instances", "semantic"] = "instances",
290    download: bool = False,
291    **kwargs
292) -> DataLoader:
293    """Get the DeepLIIF dataloader for nucleus segmentation and classification in IHC images.
294
295    Args:
296        path: The folder where the function stores the data.
297        batch_size: The batch size for training.
298        patch_shape: The patch shape to use for training.
299        split: The split of the dataset. Either 'train', 'val' or 'test'.
300        tissue: The tissue to use. See `TISSUES` for the valid choices. By default this uses
301            every tissue that the split holds. Note that 'breast' has no test split.
302        modality: The image modality to use as input. See `MODALITIES` for the valid choices.
303        label_choice: The choice of labels. Either 'instances' for the nucleus instances, or
304            'semantic' for the classification into background, negative cells and positive cells.
305        download: Whether to download the data if it is not present.
306        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
307
308    Returns:
309        The DataLoader.
310    """
311    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
312    dataset = get_deepliif_dataset(path, patch_shape, split, tissue, modality, label_choice, download, **ds_kwargs)
313    return torch_em.get_data_loader(dataset, batch_size=batch_size, **loader_kwargs)