torch_em.data.datasets.histopathology.nusec

  1"""The NuSeC dataset contains nucleus annotations for H&E stained images of invasive breast carcinoma.
  2
  3The dataset holds 100 images of 1024 x 1024 pixels at 40x magnification from 25 patients of Ankara University,
  475 for training and 25 for testing. The nuclei were outlined by hand in QuPath and checked by pathologists.
  5NOTE: The annotations are not exhaustive. Many nuclei, most of all the stromal nuclei, carry no annotation
  6and become background in the label images. Take this into account when you train on this data, and when you
  7measure the recall of a model on it.
  8
  9The data is located at https://doi.org/10.6084/m9.figshare.29525111 (CC BY 4.0), together with the
 10MiDeSeC mitosis dataset of the same group, which this module does not use.
 11This dataset is from the publication https://doi.org/10.48550/arXiv.2507.14272.
 12Please cite it if you use this dataset for your research.
 13"""
 14
 15import os
 16import zipfile
 17from tqdm import tqdm
 18from glob import glob
 19from natsort import natsorted
 20from typing import Union, Tuple, List, Literal
 21
 22import numpy as np
 23import imageio.v3 as imageio
 24
 25from torch.utils.data import Dataset, DataLoader
 26
 27import torch_em
 28
 29from .. import util
 30
 31
 32URL = "https://ndownloader.figshare.com/files/56111276"
 33CHECKSUM = "71102cec12a74f643d727b80e0a24ff12b1446a5fe3e0a94905617b28432ddf3"
 34ZIP_ROOT = "Ankara university datasets/NuSeC"
 35
 36
 37def _create_h5_files(path, zip_path):
 38    import h5py
 39
 40    with zipfile.ZipFile(zip_path) as archive:
 41        for split in ("train", "test"):
 42            out_dir = os.path.join(path, "preprocessed", split)
 43            os.makedirs(out_dir, exist_ok=True)
 44            prefix = f"{ZIP_ROOT}/{split} nuclei/"
 45            names = [n for n in archive.namelist() if n.startswith(prefix) and not n.endswith("/")]
 46            for name in tqdm(names, desc=f"Preprocess NuSeC {split} images"):
 47                fname = os.path.basename(name)
 48                image = imageio.imread(archive.read(name), extension=".tif")[..., :3]
 49                labels = imageio.imread(archive.read(f"{ZIP_ROOT}/mask of {split} nuclei/{fname}"), extension=".tif")
 50                # The masks are stored as float, but hold integer instance ids.
 51                assert np.array_equal(labels, np.round(labels)), fname
 52                labels = labels.astype("uint32")
 53                assert image.shape[:2] == labels.shape, (fname, image.shape, labels.shape)
 54
 55                with h5py.File(os.path.join(out_dir, f"{os.path.splitext(fname)[0]}.h5"), "w") as f:
 56                    f.create_dataset("raw", data=image.transpose(2, 0, 1), compression="gzip")
 57                    f.create_dataset("labels/instances", data=labels, compression="gzip")
 58
 59
 60def get_nusec_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 61    """Download the NuSeC dataset.
 62
 63    Args:
 64        path: Filepath to a folder where the downloaded data is saved.
 65        download: Whether to download the data if it is not present.
 66
 67    Returns:
 68        Filepath where the preprocessed data is stored.
 69    """
 70    data_dir = os.path.join(path, "preprocessed")
 71    if os.path.exists(data_dir):
 72        return data_dir
 73
 74    os.makedirs(path, exist_ok=True)
 75    zip_path = os.path.join(path, "nusec_midesec.zip")
 76    util.download_source(zip_path, URL, download, checksum=CHECKSUM)
 77    _create_h5_files(path, zip_path)
 78    return data_dir
 79
 80
 81def get_nusec_paths(
 82    path: Union[os.PathLike, str], split: Literal["train", "test"], download: bool = False
 83) -> List[str]:
 84    """Get paths to the NuSeC data.
 85
 86    Args:
 87        path: Filepath to a folder where the downloaded data is saved.
 88        split: The choice of data split. Either 'train' or 'test'.
 89        download: Whether to download the data if it is not present.
 90
 91    Returns:
 92        List of filepaths for the input data.
 93    """
 94    assert split in ("train", "test"), split
 95    data_dir = get_nusec_data(path, download)
 96    return natsorted(glob(os.path.join(data_dir, split, "*.h5")))
 97
 98
 99def get_nusec_dataset(
100    path: Union[os.PathLike, str],
101    patch_shape: Tuple[int, int],
102    split: Literal["train", "test"],
103    download: bool = False,
104    **kwargs
105) -> Dataset:
106    """Get the NuSeC dataset for nucleus segmentation.
107
108    Args:
109        path: Filepath to a folder where the downloaded data is saved.
110        patch_shape: The patch shape to use for training.
111        split: The choice of data split. Either 'train' or 'test'.
112        download: Whether to download the data if it is not present.
113        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
114
115    Returns:
116        The segmentation dataset.
117    """
118    paths = get_nusec_paths(path, split, download)
119    return torch_em.default_segmentation_dataset(
120        raw_paths=paths,
121        raw_key="raw",
122        label_paths=paths,
123        label_key="labels/instances",
124        patch_shape=patch_shape,
125        ndim=2,
126        with_channels=True,
127        **kwargs
128    )
129
130
131def get_nusec_loader(
132    path: Union[os.PathLike, str],
133    batch_size: int,
134    patch_shape: Tuple[int, int],
135    split: Literal["train", "test"],
136    download: bool = False,
137    **kwargs
138) -> DataLoader:
139    """Get the NuSeC dataloader for nucleus segmentation.
140
141    Args:
142        path: Filepath to a folder where the downloaded data is saved.
143        batch_size: The batch size for training.
144        patch_shape: The patch shape to use for training.
145        split: The choice of data split. Either 'train' or 'test'.
146        download: Whether to download the data if it is not present.
147        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
148
149    Returns:
150        The DataLoader.
151    """
152    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
153    dataset = get_nusec_dataset(path, patch_shape, split, download, **ds_kwargs)
154    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)