torch_em.data.datasets.histopathology.lizard_mitosis

  1"""This dataset contains annotations for nucleus instance segmentation and classification
  2in H&E stained colon histopathology images, extending the Lizard dataset with an additional
  3mitosis class. It provides two subsets: 'lizard', a modified version of the Lizard dataset
  4with mitosis annotations added, and 'mitosis', a dedicated mitosis train / validation / test
  5split (both originally released together as the 'lizard_mitosis' and 'mitosis_ds' resources).
  6
  7This dataset is from the publication https://proceedings.mlr.press/v250/baumann24a.html.
  8Please cite it if you use this dataset for your research.
  9"""
 10
 11import os
 12from tqdm import tqdm
 13from typing import Tuple, Union, Literal
 14
 15import numpy as np
 16
 17from torch.utils.data import Dataset, DataLoader
 18
 19import torch_em
 20
 21from .. import util
 22
 23
 24URLS = {
 25    "lizard": "https://zenodo.org/records/10636591/files/lizard_mitosis.zip",
 26    "mitosis": "https://zenodo.org/records/10636591/files/mitosis_ds.zip",
 27}
 28CHECKSUMS = {
 29    "lizard": "5859d738891f4620914a3fab317b3800fde8a5cd5cb4edf8003fd0efdc255ab3",
 30    "mitosis": "4b2b5d49e52611d5937f94e6cc06d8c642a0323e2250693ef465146382bde736",
 31}
 32SPLIT_FILES = {
 33    "lizard": {
 34        "train": ("fold_0/train_img.npy", "fold_0/train_lab.npy"),
 35        "val": ("fold_0/valid_img.npy", "fold_0/valid_lab.npy"),
 36        "test": ("test_images.npy", "test_labels.npy"),
 37    },
 38    "mitosis": {
 39        "train": ("train_full_img.npy", "train_full_lab.npy"),
 40        "val": ("valid_full_img.npy", "valid_full_lab.npy"),
 41        "test": ("test_ds/test_img.npy", "test_ds/test_lab.npy"),
 42    },
 43}
 44
 45
 46def _extract_split(data_dir, subset, split):
 47    import h5py
 48
 49    out_path = os.path.join(data_dir, f"{split}.h5")
 50    if os.path.exists(out_path):
 51        return out_path
 52
 53    img_file, lab_file = SPLIT_FILES[subset][split]
 54    images = np.load(os.path.join(data_dir, img_file), mmap_mode="r")
 55    labels = np.load(os.path.join(data_dir, lab_file), mmap_mode="r")
 56    assert images.shape[0] == labels.shape[0], (images.shape, labels.shape)
 57
 58    n_samples = images.shape[0]
 59    tmp_path = f"{out_path}.incomplete"
 60    with h5py.File(tmp_path, "a") as f:
 61        raw = f.create_dataset("raw", shape=(3, n_samples) + images.shape[1:3], dtype=images.dtype)
 62        instances = f.create_dataset("labels/instances", shape=(n_samples,) + images.shape[1:3], dtype=labels.dtype)
 63        semantic = f.create_dataset("labels/semantic", shape=(n_samples,) + images.shape[1:3], dtype=labels.dtype)
 64        for i in tqdm(range(n_samples), desc=f"Extract {subset} '{split}' data"):
 65            raw[:, i] = images[i].transpose(2, 0, 1)
 66            instances[i] = labels[i, ..., 0]
 67            semantic[i] = labels[i, ..., 1]
 68
 69    os.replace(tmp_path, out_path)
 70    return out_path
 71
 72
 73def get_lizard_mitosis_data(
 74    path: Union[os.PathLike, str], subset: Literal["lizard", "mitosis"], download: bool = False
 75) -> str:
 76    """Download the lizard-mitosis dataset for nucleus segmentation and classification.
 77
 78    Args:
 79        path: Filepath to a folder where the downloaded data will be saved.
 80        subset: The choice of data subset. Either 'lizard' (modified Lizard dataset with an
 81            added mitosis class) or 'mitosis' (dedicated mitosis dataset).
 82        download: Whether to download the data if it is not present.
 83
 84    Returns:
 85        The filepath to the extracted data directory.
 86    """
 87    if subset not in URLS:
 88        raise ValueError(f"'{subset}' is not a valid subset.")
 89
 90    data_dir = os.path.join(path, subset)
 91    if os.path.exists(data_dir):
 92        return data_dir
 93
 94    os.makedirs(path, exist_ok=True)
 95    zip_path = os.path.join(path, f"{subset}.zip")
 96    util.download_source(zip_path, URLS[subset], download, checksum=CHECKSUMS[subset])
 97    util.unzip(zip_path, path)
 98
 99    extracted_dir = os.path.join(path, "lizard_mitosis" if subset == "lizard" else "mitosis_ds")
100    os.rename(extracted_dir, data_dir)
101
102    return data_dir
103
104
105def get_lizard_mitosis_paths(
106    path: Union[os.PathLike, str],
107    subset: Literal["lizard", "mitosis"],
108    split: Literal["train", "val", "test"],
109    download: bool = False,
110) -> str:
111    """Get paths to the lizard-mitosis data.
112
113    Args:
114        path: Filepath to a folder where the downloaded data will be saved.
115        subset: The choice of data subset.
116        split: The choice of data split.
117        download: Whether to download the data if it is not present.
118
119    Returns:
120        Filepath to the stored data.
121    """
122    if split not in SPLIT_FILES[subset]:
123        raise ValueError(f"'{split}' is not a valid split.")
124
125    data_dir = get_lizard_mitosis_data(path, subset, download)
126    return _extract_split(data_dir, subset, split)
127
128
129def get_lizard_mitosis_dataset(
130    path: Union[os.PathLike, str],
131    patch_shape: Tuple[int, int],
132    subset: Literal["lizard", "mitosis"],
133    split: Literal["train", "val", "test"],
134    label_choice: Literal["instances", "semantic"] = "instances",
135    resize_inputs: bool = False,
136    download: bool = False,
137    **kwargs,
138) -> Dataset:
139    """Get the lizard-mitosis dataset for nucleus segmentation and classification.
140
141    Args:
142        path: Filepath to a folder where the downloaded data will be saved.
143        patch_shape: The patch shape to use for training.
144        subset: The choice of data subset.
145        split: The choice of data split.
146        label_choice: The choice of label type, either instance or semantic (class) labels.
147        resize_inputs: Whether to resize the input images.
148        download: Whether to download the data if it is not present.
149        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
150
151    Returns:
152        The segmentation dataset.
153    """
154    data_path = get_lizard_mitosis_paths(path, subset, split, download)
155
156    if resize_inputs:
157        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": True}
158        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
159            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
160        )
161
162    return torch_em.default_segmentation_dataset(
163        raw_paths=data_path,
164        raw_key="raw",
165        label_paths=data_path,
166        label_key=f"labels/{label_choice}",
167        patch_shape=patch_shape,
168        ndim=2,
169        with_channels=True,
170        **kwargs,
171    )
172
173
174def get_lizard_mitosis_loader(
175    path: Union[os.PathLike, str],
176    batch_size: int,
177    patch_shape: Tuple[int, int],
178    subset: Literal["lizard", "mitosis"],
179    split: Literal["train", "val", "test"],
180    label_choice: Literal["instances", "semantic"] = "instances",
181    resize_inputs: bool = False,
182    download: bool = False,
183    **kwargs,
184) -> DataLoader:
185    """Get the lizard-mitosis dataloader for nucleus segmentation and classification.
186
187    Args:
188        path: Filepath to a folder where the downloaded data will be saved.
189        batch_size: The batch size for training.
190        patch_shape: The patch shape to use for training.
191        subset: The choice of data subset.
192        split: The choice of data split.
193        label_choice: The choice of label type, either instance or semantic (class) labels.
194        resize_inputs: Whether to resize the inputs.
195        download: Whether to download the data if it is not present.
196        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
197
198    Returns:
199        The DataLoader.
200    """
201    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
202    dataset = get_lizard_mitosis_dataset(
203        path, patch_shape, subset, split, label_choice, resize_inputs, download, **ds_kwargs
204    )
205    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)