torch_em.data.datasets.histopathology.amonuseg

  1"""The AMONuSeg dataset contains nucleus annotations for H&E stained images from Nigeria.
  2
  3The dataset holds 48 images of 1280 x 960 pixels, taken with a microscope camera, with 12 images each of breast,
  4cervical, inguinal (lymph node) and skin tissue. Two experts annotated the nuclei: 'exp1' and 'exp2'.
  5Both annotators corrected draft masks from Fiji by hand, and four pathologists reviewed and validated them.
  6NOTE: The masks are semantic (nucleus vs. background), not instance masks. This module creates the instances as
  7connected components, so touching nuclei merge into one instance.
  8NOTE: Some mask files are stored with lossy compression artifacts, so this module binarizes all masks at half
  9the intensity range before it labels them.
 10
 11The dataset is located at https://github.com/zerouaoui/AMONUSEG (MIT license).
 12This dataset is from the publication https://doi.org/10.1007/978-3-031-72114-4_10.
 13Please cite it if you use this dataset for your research.
 14"""
 15
 16import os
 17from glob import glob
 18from tqdm import tqdm
 19from natsort import natsorted
 20from typing import Union, Tuple, List, Literal, Optional, Sequence
 21
 22import imageio.v3 as imageio
 23
 24from torch.utils.data import Dataset, DataLoader
 25
 26import torch_em
 27
 28from .. import util
 29
 30
 31# The commit of the repository that this module downloads, so that the data does not change.
 32COMMIT = "5614bf5c0aff11b3bf66420af39502d734e1c857"
 33RAW_URL = f"https://raw.githubusercontent.com/zerouaoui/AMONUSEG/{COMMIT}"
 34TREE_URL = f"https://api.github.com/repos/zerouaoui/AMONUSEG/git/trees/{COMMIT}?recursive=1"
 35ORGANS = ("breast", "cervical", "inguinal", "skin")
 36# The folders of the images and of the masks of both experts, per organ. The repository spells them differently.
 37FOLDERS = {
 38    "breast": ("Breast", "Breast Masks", "Brest Masks E2"),
 39    "cervical": ("Cervical", "Cervical Masks", "Cervical Masks E2"),
 40    "inguinal": ("Inguinal", "Inguinal Masks", "Inguinal Masks E2"),
 41    "skin": ("Skin", "Skin Masks", "Skin Masks E2"),
 42}
 43
 44
 45def _download_files(path, download):
 46    import json
 47    from urllib.parse import quote
 48
 49    tree_path = os.path.join(path, "tree.json")
 50    util.download_source(tree_path, TREE_URL, download, checksum=None)
 51    with open(tree_path) as f:
 52        files = [x["path"] for x in json.load(f)["tree"] if x["type"] == "blob" and x["path"].startswith("AMONuSeg/")]
 53    for name in files:
 54        local = os.path.join(path, name)
 55        os.makedirs(os.path.dirname(local), exist_ok=True)
 56        util.download_source(local, f"{RAW_URL}/{quote(name)}", download, checksum=None)
 57
 58
 59def _find_mask(mask_dir, stem):
 60    # The mask of image '<stem>.jpg' is '<stem>.png' or '<stem>-Mask-EXP1.png'.
 61    candidates = glob(os.path.join(mask_dir, f"{stem}.*")) + glob(os.path.join(mask_dir, f"{stem}-*"))
 62    assert len(candidates) == 1, (mask_dir, stem, candidates)
 63    return candidates[0]
 64
 65
 66def _to_instances(mask_path):
 67    from skimage.measure import label
 68
 69    mask = imageio.imread(mask_path)
 70    mask = mask[..., 0] if mask.ndim == 3 else mask
 71    return label(mask > (mask.max() / 2)).astype("uint32")
 72
 73
 74def _create_h5_files(path):
 75    import h5py
 76
 77    for organ in ORGANS:
 78        image_dir, exp1_dir, exp2_dir = (os.path.join(path, "AMONuSeg", folder) for folder in FOLDERS[organ])
 79        out_dir = os.path.join(path, "preprocessed", organ)
 80        os.makedirs(out_dir, exist_ok=True)
 81        for image_path in tqdm(natsorted(glob(os.path.join(image_dir, "*.jpg"))), desc=f"Preprocess {organ} images"):
 82            stem = os.path.splitext(os.path.basename(image_path))[0]
 83            image = imageio.imread(image_path)[..., :3]
 84            exp1 = _to_instances(_find_mask(exp1_dir, stem))
 85            exp2 = _to_instances(_find_mask(exp2_dir, stem))
 86            assert image.shape[:2] == exp1.shape == exp2.shape, stem
 87
 88            with h5py.File(os.path.join(out_dir, f"{stem}.h5"), "w") as f:
 89                f.create_dataset("raw", data=image.transpose(2, 0, 1), compression="gzip")
 90                f.create_dataset("labels/instances/exp1", data=exp1, compression="gzip")
 91                f.create_dataset("labels/instances/exp2", data=exp2, compression="gzip")
 92
 93
 94def get_amonuseg_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 95    """Download the AMONuSeg dataset.
 96
 97    Args:
 98        path: Filepath to a folder where the downloaded data is saved.
 99        download: Whether to download the data if it is not present.
100
101    Returns:
102        Filepath where the preprocessed data is stored.
103    """
104    data_dir = os.path.join(path, "preprocessed")
105    if os.path.exists(data_dir):
106        return data_dir
107
108    os.makedirs(path, exist_ok=True)
109    _download_files(path, download)
110    _create_h5_files(path)
111    return data_dir
112
113
114def get_amonuseg_paths(
115    path: Union[os.PathLike, str], organs: Optional[Sequence[str]] = None, download: bool = False
116) -> List[str]:
117    """Get paths to the AMONuSeg data.
118
119    Args:
120        path: Filepath to a folder where the downloaded data is saved.
121        organs: The organs, out of 'breast', 'cervical', 'inguinal' and 'skin'. By default, all organs are returned.
122        download: Whether to download the data if it is not present.
123
124    Returns:
125        List of filepaths for the input data.
126    """
127    data_dir = get_amonuseg_data(path, download)
128    organs = ORGANS if organs is None else organs
129    assert all(organ in ORGANS for organ in organs), organs
130    return [p for organ in organs for p in natsorted(glob(os.path.join(data_dir, organ, "*.h5")))]
131
132
133def get_amonuseg_dataset(
134    path: Union[os.PathLike, str],
135    patch_shape: Tuple[int, int],
136    organs: Optional[Sequence[str]] = None,
137    annotator: Literal["exp1", "exp2"] = "exp1",
138    download: bool = False,
139    **kwargs
140) -> Dataset:
141    """Get the AMONuSeg dataset for nucleus segmentation.
142
143    Args:
144        path: Filepath to a folder where the downloaded data is saved.
145        patch_shape: The patch shape to use for training.
146        organs: The organs, out of 'breast', 'cervical', 'inguinal' and 'skin'. By default, all organs are used.
147        annotator: The annotations of the first ('exp1') or the second ('exp2') expert.
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    assert annotator in ("exp1", "exp2"), annotator
155    paths = get_amonuseg_paths(path, organs, download)
156    return torch_em.default_segmentation_dataset(
157        raw_paths=paths,
158        raw_key="raw",
159        label_paths=paths,
160        label_key=f"labels/instances/{annotator}",
161        patch_shape=patch_shape,
162        ndim=2,
163        with_channels=True,
164        **kwargs
165    )
166
167
168def get_amonuseg_loader(
169    path: Union[os.PathLike, str],
170    batch_size: int,
171    patch_shape: Tuple[int, int],
172    organs: Optional[Sequence[str]] = None,
173    annotator: Literal["exp1", "exp2"] = "exp1",
174    download: bool = False,
175    **kwargs
176) -> DataLoader:
177    """Get the AMONuSeg dataloader for nucleus segmentation.
178
179    Args:
180        path: Filepath to a folder where the downloaded data is saved.
181        batch_size: The batch size for training.
182        patch_shape: The patch shape to use for training.
183        organs: The organs, out of 'breast', 'cervical', 'inguinal' and 'skin'. By default, all organs are used.
184        annotator: The annotations of the first ('exp1') or the second ('exp2') expert.
185        download: Whether to download the data if it is not present.
186        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
187
188    Returns:
189        The DataLoader.
190    """
191    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
192    dataset = get_amonuseg_dataset(path, patch_shape, organs, annotator, download, **ds_kwargs)
193    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)