torch_em.data.datasets.light_microscopy.ifnuclei

  1"""The IFNuclei dataset contains annotations for nucleus segmentation
  2of immuno and DAPI stained fluorescence images.
  3
  4It contains 79 expert-annotated immunofluorescence and DAPI images with 7813 nuclei from normal and
  5cancer samples. Pass `split` to use the train and test split of
  6https://github.com/kreshuklab/model_ranking, which holds 42 and 37 images.
  7
  8This dataset is from the publication https://doi.org/10.1038/s41597-020-00608-w.
  9Please cite it if you use this dataset in your research.
 10"""
 11
 12import os
 13from glob import glob
 14from natsort import natsorted
 15from typing import List, Literal, Optional, Tuple, Union
 16
 17from torch.utils.data import Dataset, DataLoader
 18
 19import torch_em
 20
 21from .. import util
 22
 23
 24URL = "https://www.ebi.ac.uk/biostudies/files/S-BSST265/dataset.zip"
 25CHECKSUM = "8285987ed4d57c46a46a55a33c1c085875ea41f429b59cde31d249741aa07ad1"
 26
 27SAMPLES = {
 28    "train": (
 29        "Ganglioneuroblastoma_0", "Ganglioneuroblastoma_1", "Ganglioneuroblastoma_2",
 30        "Ganglioneuroblastoma_3", "Neuroblastoma_0", "Neuroblastoma_1",
 31        "Neuroblastoma_10", "Neuroblastoma_11", "Neuroblastoma_2",
 32        "Neuroblastoma_3", "Neuroblastoma_4", "Neuroblastoma_5",
 33        "Neuroblastoma_6", "Neuroblastoma_7", "Neuroblastoma_8",
 34        "Neuroblastoma_9", "normal_0", "normal_1",
 35        "normal_10", "normal_11", "normal_12",
 36        "normal_13", "normal_14", "normal_15",
 37        "normal_16", "normal_17", "normal_18",
 38        "normal_19", "normal_2", "normal_20",
 39        "normal_21", "normal_22", "normal_23",
 40        "normal_24", "normal_25", "normal_3",
 41        "normal_4", "normal_5", "normal_6",
 42        "normal_7", "normal_8", "normal_9",
 43    ),
 44    "test": (
 45        "Ganglioneuroblastoma_10", "Ganglioneuroblastoma_4", "Ganglioneuroblastoma_6",
 46        "Ganglioneuroblastoma_7", "Ganglioneuroblastoma_8", "Ganglioneuroblastoma_9",
 47        "Neuroblastoma_12", "Neuroblastoma_13", "Neuroblastoma_14",
 48        "Neuroblastoma_15", "Neuroblastoma_16", "Neuroblastoma_17",
 49        "normal_26", "normal_27", "normal_28",
 50        "normal_29", "normal_30", "normal_31",
 51        "normal_32", "normal_33", "normal_34",
 52        "normal_35", "normal_36", "normal_37",
 53        "normal_38", "normal_39", "normal_40",
 54        "otherspecimen_0", "otherspecimen_1", "otherspecimen_2",
 55        "otherspecimen_3", "otherspecimen_4", "otherspecimen_5",
 56        "otherspecimen_6", "otherspecimen_7", "otherspecimen_8",
 57        "otherspecimen_9",
 58    ),
 59}
 60
 61
 62def _select_grayscale(raw):
 63    if raw.ndim == 2:
 64        return raw
 65    if raw.ndim == 3 and raw.shape[0] == 3:
 66        return raw[0]
 67    raise ValueError(f"Expected a grayscale or RGB IFNuclei image, got shape {raw.shape}.")
 68
 69
 70def get_ifnuclei_data(path: Union[os.PathLike, str], download: bool = False):
 71    """Download the IFNuclei dataset for nucleus segmentation.
 72
 73    Args:
 74        path: Filepath to a folder where the downloaded data will be saved.
 75        download: Whether to download the data if it is not present.
 76    """
 77    data_dir = os.path.join(path, "rawimages")
 78    if os.path.exists(data_dir):
 79        return
 80
 81    os.makedirs(path, exist_ok=True)
 82
 83    zip_path = os.path.join(path, "dataset.zip")
 84    util.download_source(path=zip_path, url=URL, download=download, checksum=CHECKSUM)
 85    util.unzip(zip_path=zip_path, dst=path)
 86
 87
 88def get_ifnuclei_paths(
 89    path: Union[os.PathLike, str],
 90    split: Optional[Literal["train", "test"]] = None,
 91    download: bool = False,
 92) -> Tuple[List[str], List[str]]:
 93    """Get paths to the IFNuclei data.
 94
 95    Args:
 96        path: Filepath to a folder where the downloaded data will be saved.
 97        split: The data split. Either 'train' or 'test', or None for all images.
 98        download: Whether to download the data if it is not present.
 99
100    Returns:
101        List of filepaths for the image data.
102        List of filepaths for the label data.
103    """
104    get_ifnuclei_data(path, download)
105
106    if split is None:
107        raw_paths = natsorted(glob(os.path.join(path, "rawimages", "*.tif")))
108        label_paths = natsorted(glob(os.path.join(path, "groundtruth", "*")))
109    else:
110        if split not in SAMPLES:
111            raise ValueError(f"'{split}' is not a valid split. Choose from {list(SAMPLES)}, or None for all images.")
112        raw_paths = [os.path.join(path, "rawimages", f"{sample}.tif") for sample in SAMPLES[split]]
113        label_paths = [os.path.join(path, "groundtruth", f"{sample}.tif") for sample in SAMPLES[split]]
114        missing = [p for p in raw_paths + label_paths if not os.path.exists(p)]
115        if missing:
116            raise RuntimeError(f"Could not find {len(missing)} files of the '{split}' split, e.g. {missing[0]}.")
117
118    assert len(raw_paths) == len(label_paths) and len(raw_paths) > 0
119
120    return raw_paths, label_paths
121
122
123def get_ifnuclei_dataset(
124    path: Union[os.PathLike, str],
125    patch_shape: Tuple[int, int],
126    split: Optional[Literal["train", "test"]] = None,
127    offsets: Optional[List[List[int]]] = None,
128    boundaries: bool = False,
129    binary: bool = False,
130    download: bool = False,
131    **kwargs
132) -> Dataset:
133    """Get the IFNuclei dataset for nucleus segmentation.
134
135    Args:
136        path: Filepath to a folder where the downloaded data will be saved.
137        patch_shape: The patch shape to use for training.
138        split: The data split. Either 'train' or 'test', or None for all images.
139        offsets: Offset values for affinity computation used as target.
140        boundaries: Whether to compute boundaries as the target.
141        binary: Whether to use a binary segmentation target.
142        download: Whether to download the data if it is not present.
143        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
144
145    Returns:
146        The segmentation dataset.
147    """
148    raw_paths, label_paths = get_ifnuclei_paths(path, split, download)
149
150    raw_transform = kwargs.pop("raw_transform", None)
151    if raw_transform is None:
152        raw_transform = torch_em.transform.get_raw_transform()
153    kwargs["raw_transform"] = torch_em.transform.Compose(_select_grayscale, raw_transform, is_multi_tensor=False)
154
155    kwargs, _ = util.add_instance_label_transform(
156        kwargs, add_binary_target=True, offsets=offsets, boundaries=boundaries, binary=binary,
157    )
158
159    return torch_em.default_segmentation_dataset(
160        raw_paths=raw_paths,
161        raw_key=None,
162        label_paths=label_paths,
163        label_key=None,
164        is_seg_dataset=False,
165        patch_shape=patch_shape,
166        **kwargs
167    )
168
169
170def get_ifnuclei_loader(
171    path: Union[os.PathLike, str],
172    batch_size: int,
173    patch_shape: Tuple[int, int],
174    split: Optional[Literal["train", "test"]] = None,
175    offsets: Optional[List[List[int]]] = None,
176    boundaries: bool = False,
177    binary: bool = False,
178    download: bool = False,
179    **kwargs
180) -> DataLoader:
181    """Get the IFNuclei dataloader for nucleus segmentation.
182
183    Args:
184        path: Filepath to a folder where the downloaded data will be saved.
185        batch_size: The batch size for training.
186        patch_shape: The patch shape to use for training.
187        split: The data split. Either 'train' or 'test', or None for all images.
188        offsets: Offset values for affinity computation used as target.
189        boundaries: Whether to compute boundaries as the target.
190        binary: Whether to use a binary segmentation target.
191        download: Whether to download the data if it is not present.
192        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
193
194    Returns:
195        The DataLoader.
196    """
197    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
198    dataset = get_ifnuclei_dataset(
199        path=path,
200        patch_shape=patch_shape,
201        split=split,
202        offsets=offsets,
203        boundaries=boundaries,
204        binary=binary,
205        download=download,
206        **ds_kwargs,
207    )
208    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)