torch_em.data.datasets.histopathology.wsss4luad

  1"""WSSS4LUAD is a histopathology dataset for weakly-supervised tissue semantic segmentation
  2in H&E-stained lung adenocarcinoma whole-slide images, with tumor epithelial tissue, tumor-
  3associated stroma, and normal tissue as the tissue classes.
  4
  5The original challenge data only provides patch-level (image-level) classification labels
  6for the training split. Pixel-level ground-truth segmentation masks are only provided for
  7the validation and test splits, which is why this loader restricts itself to these two splits.
  8Note that 10 of the 80 test patches have no mask in the mirrored data (their ground-truth
  9was withheld for the original challenge leaderboard); these are skipped automatically.
 10
 11The data is mirrored on Hugging Face at https://huggingface.co/datasets/Angelou0516/WSSS4LUAD-v2
 12(the original challenge data at https://wsss4luad.grand-challenge.org/WSSS4LUAD/ requires
 13requesting access via email, as the challenge has closed). The dataset is licensed under
 14CC BY 4.0. This dataset is from the publication https://doi.org/10.48550/arXiv.2204.06455.
 15Please cite it if you use this dataset in your research.
 16"""
 17
 18import os
 19from glob import glob
 20from typing import List, Literal, Tuple, Union
 21
 22from torch.utils.data import Dataset, DataLoader
 23
 24import torch_em
 25
 26from .. import util
 27
 28
 29HF_REPO = "Angelou0516/WSSS4LUAD-v2"
 30
 31SPLITS = {"val": "validation.parquet", "test": "test.parquet"}
 32
 33
 34def _extract_split(path, split):
 35    import io
 36    import numpy as np
 37    import pyarrow.parquet as pq
 38    import imageio.v3 as imageio
 39    from PIL import Image
 40    from tqdm import tqdm
 41
 42    image_dir = os.path.join(path, "images", split)
 43    mask_dir = os.path.join(path, "masks", split)
 44    if os.path.exists(image_dir) and os.path.exists(mask_dir):
 45        image_paths = sorted(glob(os.path.join(image_dir, "*.png")))
 46        if len(image_paths) > 0:
 47            return
 48
 49    os.makedirs(image_dir, exist_ok=True)
 50    os.makedirs(mask_dir, exist_ok=True)
 51
 52    parquet_path = os.path.join(path, SPLITS[split])
 53    table = pq.read_table(parquet_path)
 54
 55    for row in tqdm(table.to_pylist(), desc=f"Extracting WSSS4LUAD '{split}' split"):
 56        # A subset of the 'test' split rows have no mask in the mirrored parquet files
 57        # (the corresponding ground-truth was withheld for the original challenge leaderboard).
 58        if row["mask"] is None:
 59            continue
 60
 61        name = os.path.splitext(row["filename"])[0]
 62
 63        image_out = os.path.join(image_dir, f"{name}.png")
 64        if not os.path.exists(image_out):
 65            image = Image.open(io.BytesIO(row["image"]["bytes"])).convert("RGB")
 66            image.save(image_out)
 67
 68        mask_out = os.path.join(mask_dir, f"{name}.tif")
 69        if not os.path.exists(mask_out):
 70            mask = np.array(Image.open(io.BytesIO(row["mask"]["bytes"])))
 71            imageio.imwrite(mask_out, mask.astype("uint8"), compression="zlib")
 72
 73
 74def get_wsss4luad_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 75    """Download the WSSS4LUAD validation and test splits (the only splits with pixel-level masks).
 76
 77    Args:
 78        path: Filepath to a folder where the downloaded data will be saved.
 79        download: Whether to download the data if it is not present.
 80
 81    Returns:
 82        The filepath to the folder where the data is stored.
 83    """
 84    os.makedirs(path, exist_ok=True)
 85
 86    missing_splits = [
 87        split for split in SPLITS
 88        if not (os.path.exists(os.path.join(path, "images", split)) and glob(
 89            os.path.join(path, "images", split, "*.png")
 90        ))
 91    ]
 92    if not missing_splits:
 93        return path
 94
 95    if not download:
 96        raise RuntimeError(f"Cannot find the data at {path}, but 'download' is set to False.")
 97
 98    try:
 99        from huggingface_hub import hf_hub_download
100    except ImportError:
101        raise ImportError("'huggingface_hub' is required to download this dataset.")
102
103    for split in missing_splits:
104        hf_hub_download(
105            repo_id=HF_REPO, filename=SPLITS[split], repo_type="dataset", local_dir=path,
106        )
107        _extract_split(path, split)
108
109    return path
110
111
112def get_wsss4luad_paths(
113    path: Union[os.PathLike, str], split: Literal["val", "test"], download: bool = False,
114) -> Tuple[List[str], List[str]]:
115    """Get paths to the WSSS4LUAD image and tissue segmentation mask data.
116
117    Args:
118        path: Filepath to a folder where the downloaded data will be saved.
119        split: The data split to use. Either 'val' or 'test'.
120        download: Whether to download the data if it is not present.
121
122    Returns:
123        List of filepaths for the image data.
124        List of filepaths for the label data.
125    """
126    assert split in SPLITS, f"'{split}' is not a valid split. Choose from {list(SPLITS.keys())}."
127    data_dir = get_wsss4luad_data(path, download)
128
129    image_paths = sorted(glob(os.path.join(data_dir, "images", split, "*.png")))
130    label_paths = sorted(glob(os.path.join(data_dir, "masks", split, "*.tif")))
131    assert len(image_paths) == len(label_paths) and len(image_paths) > 0
132
133    return image_paths, label_paths
134
135
136def get_wsss4luad_dataset(
137    path: Union[os.PathLike, str],
138    patch_shape: Tuple[int, int],
139    split: Literal["val", "test"],
140    resize_inputs: bool = False,
141    download: bool = False,
142    **kwargs,
143) -> Dataset:
144    """Get the WSSS4LUAD dataset for tissue semantic segmentation.
145
146    The masks use label 0 for tumor epithelial tissue, 1 for tumor-associated stroma,
147    2 for normal tissue, and 3 for background / excluded pixels (e.g. white alveolar
148    space), which should typically not be used for computing losses or metrics.
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 'val' or 'test'.
154        resize_inputs: Whether to resize the inputs.
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    image_paths, label_paths = get_wsss4luad_paths(path, split, download)
162
163    if resize_inputs:
164        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": True}
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=image_paths,
171        raw_key=None,
172        label_paths=label_paths,
173        label_key=None,
174        patch_shape=patch_shape,
175        is_seg_dataset=False,
176        ndim=2,
177        with_channels=True,
178        **kwargs,
179    )
180
181
182def get_wsss4luad_loader(
183    path: Union[os.PathLike, str],
184    batch_size: int,
185    patch_shape: Tuple[int, int],
186    split: Literal["val", "test"],
187    resize_inputs: bool = False,
188    download: bool = False,
189    **kwargs,
190) -> DataLoader:
191    """Get the WSSS4LUAD dataloader for tissue semantic segmentation.
192
193    Args:
194        path: Filepath to a folder where the downloaded data will be saved.
195        batch_size: The batch size for training.
196        patch_shape: The patch shape to use for training.
197        split: The data split to use. Either 'val' or 'test'.
198        resize_inputs: Whether to resize the inputs.
199        download: Whether to download the data if it is not present.
200        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or the PyTorch DataLoader.
201
202    Returns:
203        The DataLoader.
204    """
205    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
206    dataset = get_wsss4luad_dataset(path, patch_shape, split, resize_inputs, download, **ds_kwargs)
207    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)