torch_em.data.datasets.histopathology.tsakiroglou

  1"""The Tsakiroglou dataset contains annotations for nucleus segmentation in DAPI-channel
  2multiplex immunofluorescence images of follicular lymphoma tissue microarray cores.
  3
  4Manual annotations have imperfect boundaries in places, and some patches show tiling
  5artifacts from stitching. Both are visible on inspection but do not prevent training use.
  6
  7The dataset is located at https://doi.org/10.17632/nb46s9trx3.1.
  8This dataset is from the publication https://doi.org/10.1007/s00262-021-02945-0.
  9Please cite it if you use this dataset for your research.
 10"""
 11
 12import os
 13from glob import glob
 14from natsort import natsorted
 15from typing import List, Literal, Tuple, Union
 16
 17from torch.utils.data import Dataset, DataLoader
 18
 19import torch_em
 20
 21from .. import util
 22
 23
 24URL = "https://data.mendeley.com/public-files/datasets/nb46s9trx3/files/e3252421-9a54-4db5-b835-1e55c184278b/file_downloaded"  # noqa
 25CHECKSUM = "1730dcaba538b03b8a1b1113242e64f59afebc9e2c4006ed559dc4251d4ade94"
 26
 27SPLIT_FOLDERS = {"train": "training_validation", "test": "testing"}
 28
 29
 30def get_tsakiroglou_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 31    """Download the Tsakiroglou dataset.
 32
 33    Args:
 34        path: Filepath to a folder where the downloaded data will be saved.
 35        download: Whether to download the data if it is not present.
 36
 37    Returns:
 38        Filepath to the folder where the data is stored.
 39    """
 40    data_dir = os.path.join(path, "nuclear_segmentation_annotations 16bit")
 41    if os.path.exists(data_dir):
 42        return data_dir
 43
 44    os.makedirs(path, exist_ok=True)
 45    zip_path = os.path.join(path, "tsakiroglou.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    return data_dir
 50
 51
 52def get_tsakiroglou_paths(
 53    path: Union[os.PathLike, str], split: Literal["train", "test"], download: bool = False
 54) -> Tuple[List[str], List[str]]:
 55    """Get paths to the Tsakiroglou data.
 56
 57    Args:
 58        path: Filepath to a folder where the downloaded data will be saved.
 59        split: The choice of data split.
 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    if split not in SPLIT_FOLDERS:
 67        raise ValueError(f"'{split}' is not a valid split choice.")
 68
 69    data_dir = get_tsakiroglou_data(path, download)
 70    split_dir = os.path.join(data_dir, SPLIT_FOLDERS[split])
 71
 72    raw_paths = natsorted(glob(os.path.join(split_dir, "DAPI_images", "*.tif")))
 73    label_paths = natsorted(glob(os.path.join(split_dir, "labels 16bit", "*.tif")))
 74
 75    assert len(raw_paths) == len(label_paths) and len(raw_paths) > 0
 76    assert all(
 77        os.path.basename(raw_path) == os.path.basename(label_path)
 78        for raw_path, label_path in zip(raw_paths, label_paths)
 79    )
 80
 81    return raw_paths, label_paths
 82
 83
 84def get_tsakiroglou_dataset(
 85    path: Union[os.PathLike, str],
 86    patch_shape: Tuple[int, int],
 87    split: Literal["train", "test"],
 88    resize_inputs: bool = False,
 89    download: bool = False,
 90    **kwargs,
 91) -> Dataset:
 92    """Get the Tsakiroglou dataset for nucleus segmentation.
 93
 94    Args:
 95        path: Filepath to a folder where the downloaded data will be saved.
 96        patch_shape: The patch shape to use for training.
 97        split: The choice of data split.
 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_tsakiroglou_paths(path, split, download)
106
107    if resize_inputs:
108        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
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        **kwargs,
121    )
122
123
124def get_tsakiroglou_loader(
125    path: Union[os.PathLike, str],
126    batch_size: int,
127    patch_shape: Tuple[int, int],
128    split: Literal["train", "test"],
129    resize_inputs: bool = False,
130    download: bool = False,
131    **kwargs,
132) -> DataLoader:
133    """Get the Tsakiroglou dataloader for nucleus segmentation.
134
135    Args:
136        path: Filepath to a folder where the downloaded data will be saved.
137        batch_size: The batch size for training.
138        patch_shape: The patch shape to use for training.
139        split: The choice of data split.
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_tsakiroglou_dataset(path, patch_shape, split, resize_inputs, download, **ds_kwargs)
149    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)