torch_em.data.datasets.histopathology.khoshdeli

  1"""The Khoshdeli dataset contains annotations for nucleus segmentation in H&E stained
  2histopathology images of brain tumor (TCGA) and breast cancer tissue.
  3
  4The dataset is located at https://doi.org/10.6084/m9.figshare.6944522.
  5This dataset is from the publication https://doi.org/10.1186/s12859-018-2285-0.
  6Please cite it if you use this dataset for your research.
  7"""
  8
  9import os
 10from glob import glob
 11from tqdm import tqdm
 12from natsort import natsorted
 13from typing import Tuple, Union, List
 14
 15import imageio.v3 as imageio
 16from bioimage_cpp.segmentation import label as connected_components
 17
 18from torch.utils.data import Dataset, DataLoader
 19
 20import torch_em
 21
 22from .. import util
 23
 24
 25URL = "https://ndownloader.figshare.com/files/12737267"
 26CHECKSUM = "929db05b0fff9139d25c8119daebe27baf6696bb6cedfc935e6e4d4be75a7620"
 27
 28
 29def get_khoshdeli_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 30    """Download the Khoshdeli dataset.
 31
 32    Args:
 33        path: Filepath to a folder where the downloaded data will be saved.
 34        download: Whether to download the data if it is not present.
 35
 36    Returns:
 37        Filepath where the data is downloaded.
 38    """
 39    data_dir = os.path.join(path, "Nuclear-Segmentation-Data")
 40    if os.path.exists(data_dir):
 41        return data_dir
 42
 43    os.makedirs(path, exist_ok=True)
 44
 45    zip_path = os.path.join(path, "khoshdeli_supplement.zip")
 46    util.download_source(path=zip_path, url=URL, download=download, checksum=CHECKSUM)
 47    util.unzip(zip_path=zip_path, dst=path)
 48
 49    tar_path = os.path.join(path, "12859_2018_2285_MOESM1_ESM", "Nuclear-Segmentation-Data.tar")
 50    util.unzip_tarfile(tar_path=tar_path, dst=path)
 51
 52    return data_dir
 53
 54
 55def get_khoshdeli_paths(path: Union[os.PathLike, str], download: bool = False) -> Tuple[List[str], List[str]]:
 56    """Get paths to the Khoshdeli data.
 57
 58    Args:
 59        path: Filepath to a folder where the downloaded data will be saved.
 60        download: Whether to download the data if it is not present.
 61
 62    Returns:
 63        List of filepaths for the image data.
 64        List of filepaths for the label data.
 65    """
 66    data_dir = get_khoshdeli_data(path, download)
 67
 68    raw_paths = natsorted(glob(os.path.join(data_dir, "Images", "*.bmp")))
 69    mask_paths = natsorted(glob(os.path.join(data_dir, "Masks", "*_MASK.bmp")))
 70    assert len(raw_paths) == len(mask_paths) and len(raw_paths) > 0
 71
 72    label_paths = []
 73    for mpath in tqdm(mask_paths, desc="Preprocessing 'khoshdeli' labels"):
 74        label_path = mpath.replace("_MASK.bmp", "_instances.tif")
 75        label_paths.append(label_path)
 76        if os.path.exists(label_path):
 77            continue
 78
 79        mask = imageio.imread(mpath) > 0
 80        label = connected_components(mask)  # run connected components to derive nucleus instances.
 81        imageio.imwrite(label_path, label, compression="zlib")
 82
 83    return raw_paths, label_paths
 84
 85
 86def get_khoshdeli_dataset(
 87    path: Union[os.PathLike, str],
 88    patch_shape: Tuple[int, int],
 89    resize_inputs: bool = False,
 90    download: bool = False,
 91    **kwargs
 92) -> Dataset:
 93    """Get the Khoshdeli dataset for nucleus segmentation.
 94
 95    Args:
 96        path: Filepath to a folder where the downloaded data will be saved.
 97        patch_shape: The patch shape to use for training.
 98        resize_inputs: Whether to resize the inputs.
 99        download: Whether to download the data if it is not present.
100        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
101
102    Returns:
103        The segmentation dataset.
104    """
105    raw_paths, label_paths = get_khoshdeli_paths(path, download)
106
107    if resize_inputs:
108        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": True}
109        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
110            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
111        )
112
113    return torch_em.default_segmentation_dataset(
114        raw_paths=raw_paths,
115        raw_key=None,
116        label_paths=label_paths,
117        label_key=None,
118        is_seg_dataset=False,
119        patch_shape=patch_shape,
120        ndim=2,
121        with_channels=True,
122        **kwargs
123    )
124
125
126def get_khoshdeli_loader(
127    path: Union[os.PathLike, str],
128    batch_size: int,
129    patch_shape: Tuple[int, int],
130    resize_inputs: bool = False,
131    download: bool = False,
132    **kwargs
133) -> DataLoader:
134    """Get the Khoshdeli dataloader for nucleus segmentation.
135
136    Args:
137        path: Filepath to a folder where the downloaded data will be saved.
138        batch_size: The batch size for training.
139        patch_shape: The patch shape to use for training.
140        resize_inputs: Whether to resize the inputs.
141        download: Whether to download the data if it is not present.
142        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
143
144    Returns:
145        The DataLoader.
146    """
147    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
148    dataset = get_khoshdeli_dataset(path, patch_shape, resize_inputs, download, **ds_kwargs)
149    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)