torch_em.data.datasets.histopathology.derma_paseg

  1"""DERMA-PASeg contains annotations for semantic segmentation of skin tissue layers in
  2brightfield whole-slide histopathology images.
  3
  4The dataset provides 32 skin biopsy sections, each imaged unstained, chemically PAS-stained,
  5and as a GAN-generated "virtually stained" counterpart, together with a semantic segmentation
  6mask for 5 classes: background and the Dermis, Epidermis, Keratin and Dermal-Epidermal
  7Junction layers. The source ships the mask as an RGB image without a color legend, so this
  8module maps its 5 colors to label ids 0-4 by sorting the RGB triplets, without claiming
  9which id corresponds to which named layer.
 10
 11NOTE: The chemically-stained image is missing for one training sample; that pair is skipped
 12when `stain="chemically_stained"`.
 13
 14The dataset is located at https://data.mendeley.com/datasets/w8vxx8yz55/1 and licensed under
 15CC BY 4.0.
 16"""
 17
 18import os
 19from glob import glob
 20from tqdm import tqdm
 21from natsort import natsorted
 22from typing import List, Literal, Tuple, Union
 23
 24import h5py
 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://data.mendeley.com/public-api/zip/w8vxx8yz55/download/1"
 36CHECKSUM = "e42604f64f2e047c8dac1a7ae23650aee91cd01dedcd65b6cec861eb0af87a62"
 37
 38# Sorted so that background (0, 0, 0) maps to label id 0.
 39LABEL_COLORS = [(0, 0, 0), (112, 48, 160), (190, 255, 0), (224, 224, 224), (255, 172, 255)]
 40
 41STAIN_FOLDERS = {"unstained": "Unstained", "chemically_stained": "C Stained", "virtually_stained": "V Stained"}
 42
 43
 44def _find_raw_path(data_dir, folder, stain, base_name):
 45    stain_dir = os.path.join(data_dir, STAIN_FOLDERS[stain], folder)
 46    if stain == "unstained":
 47        candidate = os.path.join(stain_dir, f"{base_name}.jpg")
 48        return candidate if os.path.exists(candidate) else None
 49    elif stain == "chemically_stained":
 50        candidate = os.path.join(stain_dir, f"{base_name}-PAS.jpg")
 51        return candidate if os.path.exists(candidate) else None
 52    else:  # The virtually stained images carry an inconsistent 'blended_final' / 'blended_test' suffix.
 53        matches = glob(os.path.join(stain_dir, f"{base_name}.blended_*.jpg"))
 54        return matches[0] if matches else None
 55
 56
 57def _create_h5_files(data_dir, split, stain):
 58    folder = "Train" if split == "train" else "Test"
 59    h5_dir = os.path.join(data_dir, "h5", stain, split)
 60    os.makedirs(h5_dir, exist_ok=True)
 61
 62    mask_paths = natsorted(glob(os.path.join(data_dir, "Masks", folder, "*.png")))
 63    for mask_path in tqdm(mask_paths, desc=f"Preprocessing {split} ({stain})"):
 64        base_name = os.path.splitext(os.path.basename(mask_path))[0]
 65        h5_path = os.path.join(h5_dir, f"{base_name}.h5")
 66        if os.path.exists(h5_path):
 67            continue
 68
 69        raw_path = _find_raw_path(data_dir, folder, stain, base_name)
 70        if raw_path is None:
 71            continue
 72
 73        raw = imageio.imread(raw_path)[..., :3]
 74        mask = imageio.imread(mask_path)[..., :3]
 75        labels = np.zeros(mask.shape[:2], dtype="uint8")
 76        for label_id, color in enumerate(LABEL_COLORS):
 77            labels[np.all(mask == color, axis=-1)] = label_id
 78
 79        with h5py.File(h5_path, "w") as f:
 80            f.create_dataset("raw", data=raw.transpose(2, 0, 1), compression="gzip")
 81            f.create_dataset("labels", data=labels, compression="gzip")
 82
 83
 84def get_derma_paseg_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 85    """Download the DERMA-PASeg dataset.
 86
 87    Args:
 88        path: Filepath to a folder where the downloaded data will be saved.
 89        download: Whether to download the data if it is not present.
 90
 91    Returns:
 92        The filepath to the data directory.
 93    """
 94    data_dir = os.path.join(path, "DERMA-PASeg", "DERMA-PASeg")
 95    if os.path.exists(data_dir):
 96        return data_dir
 97
 98    os.makedirs(path, exist_ok=True)
 99    zip_path = os.path.join(path, "derma_paseg.zip")
100    util.download_source(path=zip_path, url=URL, download=download, checksum=CHECKSUM)
101    util.unzip(zip_path=zip_path, dst=path)
102
103    return data_dir
104
105
106def get_derma_paseg_paths(
107    path: Union[os.PathLike, str],
108    split: Literal["train", "test"],
109    stain: Literal["unstained", "chemically_stained", "virtually_stained"] = "unstained",
110    download: bool = False,
111) -> List[str]:
112    """Get paths to the DERMA-PASeg data.
113
114    Args:
115        path: Filepath to a folder where the downloaded data will be saved.
116        split: The data split to use. Either 'train' or 'test'.
117        stain: The image variant to use as raw data. One of 'unstained', 'chemically_stained'
118            or 'virtually_stained'.
119        download: Whether to download the data if it is not present.
120
121    Returns:
122        List of filepaths for the h5 data.
123    """
124    if split not in ("train", "test"):
125        raise ValueError(f"'{split}' is not a valid split. Choose from 'train' or 'test'.")
126    if stain not in STAIN_FOLDERS:
127        raise ValueError(f"'{stain}' is not a valid stain. Choose from {list(STAIN_FOLDERS.keys())}.")
128
129    data_dir = get_derma_paseg_data(path, download)
130    _create_h5_files(data_dir, split, stain)
131
132    h5_paths = natsorted(glob(os.path.join(data_dir, "h5", stain, split, "*.h5")))
133    if len(h5_paths) == 0:
134        raise RuntimeError(f"No data found for split '{split}' and stain '{stain}'. Check the dataset at {data_dir}.")
135
136    return h5_paths
137
138
139def get_derma_paseg_dataset(
140    path: Union[os.PathLike, str],
141    patch_shape: Tuple[int, int],
142    split: Literal["train", "test"],
143    stain: Literal["unstained", "chemically_stained", "virtually_stained"] = "unstained",
144    resize_inputs: bool = False,
145    download: bool = False,
146    **kwargs,
147) -> Dataset:
148    """Get the DERMA-PASeg dataset for skin tissue layer segmentation.
149
150    Args:
151        path: Filepath to a folder where the downloaded data will be saved.
152        patch_shape: The patch shape to use for training.
153        split: The data split to use. Either 'train' or 'test'.
154        stain: The image variant to use as raw data. One of 'unstained', 'chemically_stained'
155            or 'virtually_stained'.
156        resize_inputs: Whether to resize the inputs.
157        download: Whether to download the data if it is not present.
158        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
159
160    Returns:
161        The segmentation dataset.
162    """
163    h5_paths = get_derma_paseg_paths(path, split, stain, download)
164
165    if resize_inputs:
166        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": True}
167        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
168            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
169        )
170
171    return torch_em.default_segmentation_dataset(
172        raw_paths=h5_paths,
173        raw_key="raw",
174        label_paths=h5_paths,
175        label_key="labels",
176        patch_shape=patch_shape,
177        with_channels=True,
178        ndim=2,
179        **kwargs,
180    )
181
182
183def get_derma_paseg_loader(
184    path: Union[os.PathLike, str],
185    batch_size: int,
186    patch_shape: Tuple[int, int],
187    split: Literal["train", "test"],
188    stain: Literal["unstained", "chemically_stained", "virtually_stained"] = "unstained",
189    resize_inputs: bool = False,
190    download: bool = False,
191    **kwargs,
192) -> DataLoader:
193    """Get the DERMA-PASeg dataloader for skin tissue layer segmentation.
194
195    Args:
196        path: Filepath to a folder where the downloaded data will be saved.
197        batch_size: The batch size for training.
198        patch_shape: The patch shape to use for training.
199        split: The data split to use. Either 'train' or 'test'.
200        stain: The image variant to use as raw data. One of 'unstained', 'chemically_stained'
201            or 'virtually_stained'.
202        resize_inputs: Whether to resize the inputs.
203        download: Whether to download the data if it is not present.
204        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
205
206    Returns:
207        The DataLoader.
208    """
209    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
210    dataset = get_derma_paseg_dataset(path, patch_shape, split, stain, resize_inputs, download, **ds_kwargs)
211    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)