torch_em.data.datasets.histopathology.glysac

  1"""The GLySAC dataset contains annotations for nuclei instance segmentation and
  2classification in H&E stained gastric cancer histopathology images.
  3
  4The dataset contains 59 image tiles of size 1000x1000 pixels with instance
  5segmentation masks and cell type annotations. Three cell classes are provided:
  6lymphocytes, epithelial cells (normal and tumor), and other cells.
  7
  8NOTE: The dataset is hosted on Google Drive and requires gdown to download.
  9Install it with: conda install -c conda-forge gdown==4.6.3
 10
 11The dataset is located at https://drive.google.com/file/d/1g1_xYFWgp3cRLKrlSwD2U5JDjooC0yHp/view
 12This dataset is from the publication https://doi.org/10.1109/jbhi.2022.3149936.
 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 List, Literal, Tuple, Union
 21
 22import h5py
 23import imageio.v3 as imageio
 24from scipy.io import loadmat
 25from torch.utils.data import Dataset, DataLoader
 26
 27import torch_em
 28
 29from .. import util
 30
 31
 32GDRIVE_ID = "1g1_xYFWgp3cRLKrlSwD2U5JDjooC0yHp"
 33URL = f"https://drive.google.com/uc?id={GDRIVE_ID}"
 34CHECKSUM = None
 35
 36
 37def _find_split_dir(data_dir: str, split: str) -> str:
 38    """Resolve the split folder. Distributions of this dataset differ in capitalization."""
 39    for folder in (split.capitalize(), split, split.upper()):
 40        candidate = os.path.join(data_dir, folder)
 41        if os.path.exists(candidate):
 42            return candidate
 43    raise RuntimeError(f"Could not find a folder for split '{split}' in '{data_dir}'.")
 44
 45
 46def _create_h5_files(data_dir: str, split: str) -> None:
 47    split_dir = _find_split_dir(data_dir, split)
 48    image_dir = os.path.join(split_dir, "Images")
 49    label_dir = os.path.join(split_dir, "Labels")
 50    h5_dir = os.path.join(data_dir, "h5", split)
 51    os.makedirs(h5_dir, exist_ok=True)
 52
 53    # Distributions of this dataset ship the images either as png or as tif.
 54    image_paths = natsorted(glob(os.path.join(image_dir, "*.png")) + glob(os.path.join(image_dir, "*.tif")))
 55    for image_path in tqdm(image_paths, desc=f"Preprocessing {split}"):
 56        fname = os.path.splitext(os.path.basename(image_path))[0]
 57        h5_path = os.path.join(h5_dir, f"{fname}.h5")
 58        if os.path.exists(h5_path):
 59            continue
 60
 61        label_path = os.path.join(label_dir, f"{fname}.mat")
 62        raw = imageio.imread(image_path)[..., :3]
 63        mat = loadmat(label_path)
 64        inst_map = mat["inst_map"].astype("int32")
 65        type_map = mat["type_map"].astype("int32")
 66
 67        with h5py.File(h5_path, "w") as f:
 68            f.create_dataset("raw", data=raw.transpose(2, 0, 1), compression="gzip")
 69            f.create_dataset("labels/instances", data=inst_map, compression="gzip")
 70            f.create_dataset("labels/semantic", data=type_map, compression="gzip")
 71
 72
 73def get_glysac_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 74    """Download the GLySAC dataset.
 75
 76    Args:
 77        path: Filepath to a folder where the downloaded data will be saved.
 78        download: Whether to download the data if it is not present.
 79
 80    Returns:
 81        The filepath to the data directory.
 82    """
 83    # Distributions of this dataset unpack either as 'glysac_dataset' or as 'GLySAC'.
 84    for folder in ("glysac_dataset", "GLySAC"):
 85        data_dir = os.path.join(path, folder)
 86        if os.path.exists(data_dir):
 87            return data_dir
 88
 89    data_dir = os.path.join(path, "glysac_dataset")
 90
 91    os.makedirs(path, exist_ok=True)
 92    zip_path = os.path.join(path, "glysac_dataset.zip")
 93    util.download_source_gdrive(path=zip_path, url=URL, download=download, checksum=CHECKSUM)
 94    util.unzip(zip_path, path)
 95
 96    return data_dir
 97
 98
 99def get_glysac_paths(
100    path: Union[os.PathLike, str],
101    split: Literal["train", "test"],
102    download: bool = False,
103) -> List[str]:
104    """Get paths to the GLySAC data.
105
106    Args:
107        path: Filepath to a folder where the downloaded data will be saved.
108        split: The data split to use. Either 'train' or 'test'.
109        download: Whether to download the data if it is not present.
110
111    Returns:
112        List of filepaths for the h5 data.
113    """
114    if split not in ("train", "test"):
115        raise ValueError(f"'{split}' is not a valid split. Choose from 'train' or 'test'.")
116
117    data_dir = get_glysac_data(path, download)
118    _create_h5_files(data_dir, split)
119
120    h5_paths = natsorted(glob(os.path.join(data_dir, "h5", split, "*.h5")))
121    if len(h5_paths) == 0:
122        raise RuntimeError(f"No data found for split '{split}'. Check the dataset at {data_dir}.")
123
124    return h5_paths
125
126
127def get_glysac_dataset(
128    path: Union[os.PathLike, str],
129    patch_shape: Tuple[int, int],
130    split: Literal["train", "test"],
131    label_choice: Literal["instances", "semantic"] = "instances",
132    download: bool = False,
133    **kwargs,
134) -> Dataset:
135    """Get the GLySAC dataset for gastric nuclei segmentation.
136
137    Args:
138        path: Filepath to a folder where the downloaded data will be saved.
139        patch_shape: The patch shape to use for training.
140        split: The data split to use. Either 'train' or 'test'.
141        label_choice: The type of labels to load. Either 'instances' for instance segmentation
142            or 'semantic' for cell type classification (4 classes: other, lymphocyte, epithelial, ambiguous).
143        download: Whether to download the data if it is not present.
144        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
145
146    Returns:
147        The segmentation dataset.
148    """
149    if label_choice not in ("instances", "semantic"):
150        raise ValueError(f"'{label_choice}' is not a valid label choice. Use 'instances' or 'semantic'.")
151
152    h5_paths = get_glysac_paths(path, split, download)
153
154    if label_choice == "instances":
155        kwargs, _ = util.add_instance_label_transform(kwargs, add_binary_target=True)
156    kwargs = util.ensure_transforms(ndim=2, **kwargs)
157
158    return torch_em.default_segmentation_dataset(
159        raw_paths=h5_paths,
160        raw_key="raw",
161        label_paths=h5_paths,
162        label_key=f"labels/{label_choice}",
163        patch_shape=patch_shape,
164        with_channels=True,
165        ndim=2,
166        **kwargs,
167    )
168
169
170def get_glysac_loader(
171    path: Union[os.PathLike, str],
172    batch_size: int,
173    patch_shape: Tuple[int, int],
174    split: Literal["train", "test"],
175    label_choice: Literal["instances", "semantic"] = "instances",
176    download: bool = False,
177    **kwargs,
178) -> DataLoader:
179    """Get the GLySAC dataloader for gastric nuclei segmentation.
180
181    Args:
182        path: Filepath to a folder where the downloaded data will be saved.
183        batch_size: The batch size for training.
184        patch_shape: The patch shape to use for training.
185        split: The data split to use. Either 'train' or 'test'.
186        label_choice: The type of labels to load. Either 'instances' for instance segmentation
187            or 'semantic' for cell type classification (4 classes: other, lymphocyte, epithelial, ambiguous).
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_glysac_dataset(path, patch_shape, split, label_choice, download, **ds_kwargs)
196    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)