torch_em.data.datasets.light_microscopy.hela_cytonuc

  1"""The HeLaCytoNuc dataset contains fluorescence images of HeLa cells with instance annotations.
  2
  3The red image channel shows the cytoplasm, the blue channel shows the nuclei and the green channel is unused.
  4The dataset is available at https://doi.org/10.14278/rodare.3001 under the CC BY 4.0 license.
  5Please cite the dataset record if you use this dataset in your research.
  6"""
  7
  8import os
  9from glob import glob
 10from operator import itemgetter
 11from typing import List, Literal, Tuple, Union
 12
 13from torch.utils.data import DataLoader, Dataset
 14
 15import torch_em
 16
 17from .. import util
 18
 19
 20BASE_URL = "https://rodare.hzdr.de/api/files/fae71336-c6d2-45b2-ae65-416c8f57b5a0"
 21URLS = {
 22    "train": f"{BASE_URL}/HeLaCytoNuc_train.zip",
 23    "val": f"{BASE_URL}/HeLaCytoNuc_validation.zip",
 24    "test": f"{BASE_URL}/HeLaCytoNuc_test.zip",
 25}
 26CHECKSUMS = {
 27    "train": "9241233246977df5b177f038257855773996ed4cdb5682f86d3a9d3f127970f3",
 28    "val": "37c4c0904db2146620d5f9b4df4c116768712dafead8c8008b38f249a034bc8b",
 29    "test": "9f50e9b5df435cafb95d3434571df5cf70e0af387f5cdeb06081fb521f5ded0e",
 30}
 31EXPECTED_SAMPLES = {"train": 1873, "val": 535, "test": 268}
 32ARCHIVE_NAMES = {"train": "train", "val": "validation", "test": "test"}
 33RAW_CHANNELS = {"cytoplasm": 0, "nuclei": 2}
 34
 35
 36def get_hela_cytonuc_data(
 37    path: Union[os.PathLike, str], split: Literal["train", "val", "test"], download: bool = False,
 38) -> str:
 39    """Download the HeLaCytoNuc data for one split.
 40
 41    Args:
 42        path: Filepath to a folder where the downloaded data will be saved.
 43        split: The data split. One of 'train', 'val', or 'test'.
 44        download: Whether to download the data if it is not present.
 45
 46    Returns:
 47        The filepath to the data for the requested split.
 48    """
 49    if split not in URLS:
 50        raise ValueError(f"'{split}' is not a valid split. Choose from {list(URLS)}.")
 51
 52    split_path = os.path.join(path, split)
 53    data_folders = [os.path.join(split_path, name) for name in ("images", "nuclei_masks", "cytoplasm_masks")]
 54    if all(os.path.exists(folder) for folder in data_folders):
 55        return split_path
 56
 57    os.makedirs(split_path, exist_ok=True)
 58    archive_name = ARCHIVE_NAMES[split]
 59    zip_path = os.path.join(path, f"HeLaCytoNuc_{archive_name}.zip")
 60    util.download_source(zip_path, URLS[split], download, CHECKSUMS[split])
 61    util.unzip(zip_path, split_path)
 62
 63    if not all(os.path.exists(folder) for folder in data_folders):
 64        raise RuntimeError(f"The downloaded archive for split '{split}' has an unexpected structure.")
 65    return split_path
 66
 67
 68def get_hela_cytonuc_paths(
 69    path: Union[os.PathLike, str],
 70    split: Literal["train", "val", "test"],
 71    label_choice: Literal["nuclei", "cytoplasm"] = "nuclei",
 72    download: bool = False,
 73) -> Tuple[List[str], List[str]]:
 74    """Get paths to the HeLaCytoNuc images and instance labels.
 75
 76    Args:
 77        path: Filepath to a folder where the downloaded data will be saved.
 78        split: The data split. One of 'train', 'val', or 'test'.
 79        label_choice: The instance annotations to load. Either 'nuclei' or 'cytoplasm'.
 80        download: Whether to download the data if it is not present.
 81
 82    Returns:
 83        The image paths and corresponding label paths.
 84    """
 85    if label_choice not in ("nuclei", "cytoplasm"):
 86        raise ValueError("The label choice must be either 'nuclei' or 'cytoplasm'.")
 87
 88    split_path = get_hela_cytonuc_data(path, split, download)
 89    image_paths = sorted(glob(os.path.join(split_path, "images", "*.tif")))
 90    expected_samples = EXPECTED_SAMPLES[split]
 91    if len(image_paths) != expected_samples:
 92        raise RuntimeError(
 93            f"Expected {expected_samples} images for split '{split}', but found {len(image_paths)}."
 94        )
 95
 96    label_folder = os.path.join(split_path, f"{label_choice}_masks")
 97    label_paths = [os.path.join(label_folder, os.path.basename(image_path)) for image_path in image_paths]
 98    missing_labels = [label_path for label_path in label_paths if not os.path.exists(label_path)]
 99    if missing_labels:
100        raise RuntimeError(f"Could not find labels for {len(missing_labels)} images in '{label_folder}'.")
101
102    return image_paths, label_paths
103
104
105def get_hela_cytonuc_dataset(
106    path: Union[os.PathLike, str],
107    patch_shape: Tuple[int, int],
108    split: Literal["train", "val", "test"],
109    raw_channel: Literal["rgb", "nuclei", "cytoplasm"] = "rgb",
110    label_choice: Literal["nuclei", "cytoplasm"] = "nuclei",
111    download: bool = False,
112    **kwargs,
113) -> Dataset:
114    """Get the HeLaCytoNuc dataset for nucleus or cytoplasm instance segmentation.
115
116    Args:
117        path: Filepath to a folder where the downloaded data will be saved.
118        patch_shape: The patch shape to use for training.
119        split: The data split. One of 'train', 'val', or 'test'.
120        raw_channel: The image channels to load. Either 'rgb', 'nuclei', or 'cytoplasm'.
121        label_choice: The instance annotations to load. Either 'nuclei' or 'cytoplasm'.
122        download: Whether to download the data if it is not present.
123        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
124
125    Returns:
126        The segmentation dataset.
127    """
128    if raw_channel not in ("rgb", "nuclei", "cytoplasm"):
129        raise ValueError("The raw channel must be 'rgb', 'nuclei', or 'cytoplasm'.")
130
131    image_paths, label_paths = get_hela_cytonuc_paths(path, split, label_choice, download)
132
133    if raw_channel != "rgb":
134        raw_transform = kwargs.pop("raw_transform", None)
135        if raw_transform is None:
136            raw_transform = torch_em.transform.get_raw_transform()
137        kwargs["raw_transform"] = torch_em.transform.Compose(
138            itemgetter(RAW_CHANNELS[raw_channel]), raw_transform, is_multi_tensor=False,
139        )
140
141    kwargs = util.update_kwargs(kwargs, "is_seg_dataset", False)
142
143    return torch_em.default_segmentation_dataset(
144        raw_paths=image_paths,
145        raw_key=None,
146        label_paths=label_paths,
147        label_key=None,
148        patch_shape=patch_shape,
149        **kwargs,
150    )
151
152
153def get_hela_cytonuc_loader(
154    path: Union[os.PathLike, str],
155    batch_size: int,
156    patch_shape: Tuple[int, int],
157    split: Literal["train", "val", "test"],
158    raw_channel: Literal["rgb", "nuclei", "cytoplasm"] = "rgb",
159    label_choice: Literal["nuclei", "cytoplasm"] = "nuclei",
160    download: bool = False,
161    **kwargs,
162) -> DataLoader:
163    """Get the HeLaCytoNuc dataloader for nucleus or cytoplasm instance segmentation.
164
165    Args:
166        path: Filepath to a folder where the downloaded data will be saved.
167        batch_size: The batch size for training.
168        patch_shape: The patch shape to use for training.
169        split: The data split. One of 'train', 'val', or 'test'.
170        raw_channel: The image channels to load. Either 'rgb', 'nuclei', or 'cytoplasm'.
171        label_choice: The instance annotations to load. Either 'nuclei' or 'cytoplasm'.
172        download: Whether to download the data if it is not present.
173        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
174
175    Returns:
176        The DataLoader.
177    """
178    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
179    dataset = get_hela_cytonuc_dataset(
180        path=path,
181        patch_shape=patch_shape,
182        split=split,
183        raw_channel=raw_channel,
184        label_choice=label_choice,
185        download=download,
186        **ds_kwargs,
187    )
188    return torch_em.get_data_loader(dataset, batch_size=batch_size, **loader_kwargs)