torch_em.data.datasets.light_microscopy.dcis_com_nuclei

  1"""The DCIS.COM nuclei dataset contains SiR-DNA fluorescence images with instance annotations.
  2
  3The dataset contains images of DCIS.COM cells acquired with spinning disk confocal microscopy.
  4It is located at https://www.ebi.ac.uk/biostudies/bioimages/studies/S-BIAD895.
  5The original dataset is available at https://doi.org/10.5281/zenodo.3715492 under the CC BY 4.0 license.
  6Please cite the dataset record if you use this dataset in your research.
  7"""
  8
  9import os
 10from glob import glob
 11from shutil import rmtree
 12from typing import List, Literal, Optional, Tuple, Union
 13
 14from torch.utils.data import DataLoader, Dataset
 15
 16import torch_em
 17
 18from .. import util
 19
 20
 21URL = "https://zenodo.org/records/3715492/files/Stardist_v2.zip?download=1"
 22CHECKSUM = "aec767afae76942b7c97e31c500284f8b5862150d8e81b57f513e66d7258c05e"
 23
 24SPLIT_FOLDERS = {
 25    "train": ("Training - Images", "Training - Masks"),
 26    "test": ("Test - Images", "Test - Masks"),
 27}
 28EXPECTED_SAMPLES = {"train": 45, "test": 2}
 29INVALID_TRAIN_IMAGES = {
 30    "cell migration R1 - Position 0_XY1562686096_Z0_T00_C1-1-image7.tif",
 31    "cell migration R1 - Position 0_XY1562686096_Z0_T00_C1-1-image14.tif",
 32}
 33
 34
 35def get_dcis_com_nuclei_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 36    """Download the DCIS.COM nuclei dataset (S-BIAD895).
 37
 38    Args:
 39        path: Filepath to a folder where the downloaded data will be saved.
 40        download: Whether to download the data if it is not present.
 41
 42    Returns:
 43        The filepath to the downloaded data.
 44    """
 45    data_dir = os.path.join(path, "Stardist")
 46    expected_folders = [os.path.join(data_dir, folder) for folders in SPLIT_FOLDERS.values() for folder in folders]
 47    if all(os.path.exists(folder) for folder in expected_folders):
 48        return data_dir
 49
 50    os.makedirs(path, exist_ok=True)
 51    zip_path = os.path.join(path, "Stardist_v2.zip")
 52    util.download_source(zip_path, URL, download, CHECKSUM)
 53    util.unzip(zip_path, path)
 54
 55    macos_dir = os.path.join(path, "__MACOSX")
 56    if os.path.exists(macos_dir):
 57        rmtree(macos_dir)
 58
 59    if not all(os.path.exists(folder) for folder in expected_folders):
 60        raise RuntimeError("The downloaded S-BIAD895 archive has an unexpected structure.")
 61    return data_dir
 62
 63
 64def get_dcis_com_nuclei_paths(
 65    path: Union[os.PathLike, str],
 66    split: Literal["train", "test"] = "train",
 67    download: bool = False,
 68) -> Tuple[List[str], List[str]]:
 69    """Get paths to the DCIS.COM nuclei images and instance labels.
 70
 71    Two training images are excluded because their source masks are exact copies of the mask for
 72    ``image1.tif`` and do not align with the corresponding images.
 73
 74    Args:
 75        path: Filepath to a folder where the downloaded data will be saved.
 76        split: The data split. Either 'train' or 'test'.
 77        download: Whether to download the data if it is not present.
 78
 79    Returns:
 80        The image paths and corresponding label paths.
 81    """
 82    if split not in SPLIT_FOLDERS:
 83        raise ValueError(f"'{split}' is not a valid split. Choose from {list(SPLIT_FOLDERS)}.")
 84
 85    data_dir = get_dcis_com_nuclei_data(path, download)
 86    image_folder, label_folder = SPLIT_FOLDERS[split]
 87    image_paths = sorted(glob(os.path.join(data_dir, image_folder, "*.tif")))
 88
 89    expected_samples = EXPECTED_SAMPLES[split]
 90    if len(image_paths) != expected_samples:
 91        raise RuntimeError(
 92            f"Expected {expected_samples} images for split '{split}', but found {len(image_paths)}."
 93        )
 94
 95    if split == "train":
 96        image_paths = [
 97            image_path for image_path in image_paths if os.path.basename(image_path) not in INVALID_TRAIN_IMAGES
 98        ]
 99
100    label_folder = os.path.join(data_dir, label_folder)
101    label_paths = [os.path.join(label_folder, os.path.basename(image_path)) for image_path in image_paths]
102    missing_labels = [label_path for label_path in label_paths if not os.path.exists(label_path)]
103    if missing_labels:
104        raise RuntimeError(f"Could not find labels for {len(missing_labels)} images in '{label_folder}'.")
105
106    return image_paths, label_paths
107
108
109def get_dcis_com_nuclei_dataset(
110    path: Union[os.PathLike, str],
111    patch_shape: Tuple[int, int],
112    split: Literal["train", "test"] = "train",
113    offsets: Optional[List[List[int]]] = None,
114    boundaries: bool = False,
115    binary: bool = False,
116    download: bool = False,
117    **kwargs,
118) -> Dataset:
119    """Get the DCIS.COM nuclei dataset for instance segmentation.
120
121    Args:
122        path: Filepath to a folder where the downloaded data will be saved.
123        patch_shape: The patch shape to use for training.
124        split: The data split. Either 'train' or 'test'.
125        offsets: Offset values for affinity computation used as target.
126        boundaries: Whether to compute boundaries as the target.
127        binary: Whether to use a binary segmentation target.
128        download: Whether to download the data if it is not present.
129        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
130
131    Returns:
132        The segmentation dataset.
133    """
134    image_paths, label_paths = get_dcis_com_nuclei_paths(path, split, download)
135
136    kwargs, _ = util.add_instance_label_transform(
137        kwargs, add_binary_target=True, offsets=offsets, boundaries=boundaries, binary=binary,
138    )
139    kwargs = util.update_kwargs(kwargs, "is_seg_dataset", False)
140
141    return torch_em.default_segmentation_dataset(
142        raw_paths=image_paths,
143        raw_key=None,
144        label_paths=label_paths,
145        label_key=None,
146        patch_shape=patch_shape,
147        **kwargs,
148    )
149
150
151def get_dcis_com_nuclei_loader(
152    path: Union[os.PathLike, str],
153    batch_size: int,
154    patch_shape: Tuple[int, int],
155    split: Literal["train", "test"] = "train",
156    offsets: Optional[List[List[int]]] = None,
157    boundaries: bool = False,
158    binary: bool = False,
159    download: bool = False,
160    **kwargs,
161) -> DataLoader:
162    """Get the DCIS.COM nuclei dataloader for instance segmentation.
163
164    Args:
165        path: Filepath to a folder where the downloaded data will be saved.
166        batch_size: The batch size for training.
167        patch_shape: The patch shape to use for training.
168        split: The data split. Either 'train' or 'test'.
169        offsets: Offset values for affinity computation used as target.
170        boundaries: Whether to compute boundaries as the target.
171        binary: Whether to use a binary segmentation target.
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_dcis_com_nuclei_dataset(
180        path=path,
181        patch_shape=patch_shape,
182        split=split,
183        offsets=offsets,
184        boundaries=boundaries,
185        binary=binary,
186        download=download,
187        **ds_kwargs,
188    )
189    return torch_em.get_data_loader(dataset, batch_size=batch_size, **loader_kwargs)