torch_em.data.datasets.light_microscopy.cardioblast_nuclei

  1"""The cardioblast nuclei dataset contains annotated time-lapse fluorescence microscopy images.
  2
  3It shows cardioblast nuclei migrating during Drosophila embryonic development to form the early heart tube.
  4Each TIFF stack contains a time series of maximum-intensity projections with tracked nucleus instance labels.
  5One raw movie has 26 unannotated trailing frames; the loader restricts this movie to its annotated prefix.
  6This loader exposes the subset and train/test split used by https://github.com/kreshuklab/model_ranking.
  7
  8The dataset is located at https://doi.org/10.6019/S-BIAD1410 and is available under the CC0 license.
  9It is from the publication https://doi.org/10.1083/jcb.202506102.
 10Please cite the dataset and publication if you use this dataset in your research.
 11"""
 12
 13import os
 14from typing import List, Literal, Optional, Tuple, Union
 15
 16from torch.utils.data import DataLoader, Dataset
 17
 18import torch_em
 19
 20from .. import util
 21
 22
 23BASE_URL = "https://www.ebi.ac.uk/biostudies/files/S-BIAD1410/cardioblast_nuclei"
 24
 25SAMPLES = {
 26    "train": (
 27        "cardioblast_nuclei_20200127_e1",
 28        "cardioblast_nuclei_20200131_e1",
 29        "cardioblast_nuclei_20200206_e1",
 30        "cardioblast_nuclei_20200206_e3",
 31        "cardioblast_nuclei_20220811_e1",
 32        "cardioblast_nuclei_20220812_e1",
 33        "cardioblast_nuclei_20220826_e2",
 34        "cardioblast_nuclei_20220828_e1",
 35        "cardioblast_nuclei_20220828_e2",
 36        "cardioblast_nuclei_20220828_e4",
 37    ),
 38    "test": (
 39        "cardioblast_nuclei_20200121_e3",
 40        "cardioblast_nuclei_20200219_e1",
 41        "cardioblast_nuclei_20220811_e2",
 42        "cardioblast_nuclei_20220825_e1",
 43        "cardioblast_nuclei_20220828_e3",
 44    ),
 45}
 46
 47RAW_CHECKSUMS = {
 48    "cardioblast_nuclei_20200127_e1": "eaedb22edcef7bfa5b8bb099ed13270b5c7f099178be6c3857341874c5582f8c",
 49    "cardioblast_nuclei_20200131_e1": "6efa5753a639ea02a627ba2d6fdbd36c13fc42b531881452305e2d769231a40b",
 50    "cardioblast_nuclei_20200206_e1": "0f347dc898e89ba1f35abd21a7473759a55891db9f79e9062a521cf865677b1f",
 51    "cardioblast_nuclei_20200206_e3": "0334aff3adec20a8df53b5f628a2523b50d9eee5e8e83797eeb429be4b22d9b3",
 52    "cardioblast_nuclei_20220811_e1": "6efec42327f5cc81c83534b5762591a36841ca1f4c806ffc2d43abd7eb0b6f38",
 53    "cardioblast_nuclei_20220812_e1": "452f228ac7e0889eb2ffa7c9a264d372af12f25bf51526380e1ff14300e6367a",
 54    "cardioblast_nuclei_20220826_e2": "9d7d5e9572b147b47d4a502cc9d06defad67ceda7574b1283d9b77c7e52fbc99",
 55    "cardioblast_nuclei_20220828_e1": "24d7e20ea7fd268064536f8bc2e3a28cc79a04276fa848e6e029940ef5fb1489",
 56    "cardioblast_nuclei_20220828_e2": "55a46cbc76e1831c30cfde0e6454168e14c126246ec0cf00e96459e3925a67b3",
 57    "cardioblast_nuclei_20220828_e4": "82c6cfde80efb1b8416fb8533dd94c73cea227274e8bb7fb3a610309ae7c0620",
 58    "cardioblast_nuclei_20200121_e3": "dcfd039fe12050c6c75b74fe315fc6cf1f7295fac24f32111e416bb0977f59e2",
 59    "cardioblast_nuclei_20200219_e1": "b4605fda7b11e091af1e431df81bb96b79e86cb63081e259822d429475f3e869",
 60    "cardioblast_nuclei_20220811_e2": "a79ff0dcf65d5cf823d1a3f29ed7751f44e08a3545010054b5d95648b8174bd8",
 61    "cardioblast_nuclei_20220825_e1": "1608f2441090c2f5366e0afbdc980e1c2576444dfbb10bc7b1c9d5ba00fedbb0",
 62    "cardioblast_nuclei_20220828_e3": "5e8f05ff702a4552495989a20df066322074f571c97797f525e7d8a85e2f5862",
 63}
 64
 65LABEL_CHECKSUMS = {
 66    "cardioblast_nuclei_20200127_e1": "2256bc19a75f7935732dd0ffab43602f5ba58167d828bc0d16c57325929c19a3",
 67    "cardioblast_nuclei_20200131_e1": "033db5065e39abbadd8d6accd83b330271fd8432fc2cbadeaf680a3bdf8a5f67",
 68    "cardioblast_nuclei_20200206_e1": "2a7bed91aa3e38542731bd5854abcc2db593dfd1a0329781be52dac39ffd8ece",
 69    "cardioblast_nuclei_20200206_e3": "6e133822ef6d1ce99b9f7ee66fac16edcb50ef01bd85971cbde3dd2fa6630cee",
 70    "cardioblast_nuclei_20220811_e1": "1d95565d1d5a925a46e0e8b58055772ed22b8821e6bf5c9d3152e0eeec1764ed",
 71    "cardioblast_nuclei_20220812_e1": "2224859e84a811210607759ef3c5c46efca07a6ab1f8948889f7e5c4cb0d4a6a",
 72    "cardioblast_nuclei_20220826_e2": "6dc737a7e0aab77fab1b0489e41fa8eea50ea95cb32866f02c2ad72ae7a78314",
 73    "cardioblast_nuclei_20220828_e1": "b0f64623fd63ee1f9fdad05fef1bc171d7fed70a470cab68feb236fc0fbff636",
 74    "cardioblast_nuclei_20220828_e2": "0182987832eba600b9bc2c31048de42d427575abdda8b35809062278f18bb27d",
 75    "cardioblast_nuclei_20220828_e4": "0f0410ec34e3536beadc7403c98ab7786f87659e31f0739505d807a905ea3292",
 76    "cardioblast_nuclei_20200121_e3": "a438457b17036de1677557d6c2d56454e36a5e950ca7d3c8e5e263f449962cd0",
 77    "cardioblast_nuclei_20200219_e1": "a783b5f3c5b460a009f632875f8c1e456903acaf4c2f2531695e8bbd736f808e",
 78    "cardioblast_nuclei_20220811_e2": "5b6dcdfbc90a0955c129706c3e8defb8003ff6bc17139b03a74a2f81945742c0",
 79    "cardioblast_nuclei_20220825_e1": "ffcbe838b2887bc8ac02ce0380b768febfcc65e5709abf7b6522481197e5f67f",
 80    "cardioblast_nuclei_20220828_e3": "bee487d987577f93c7061d884d0fc832bba90854f348b0734208ec631f30b06d",
 81}
 82
 83# This movie contains 147 raw frames, but only the first 121 frames are annotated.
 84ANNOTATED_FRAMES = {"cardioblast_nuclei_20220812_e1": 121}
 85
 86
 87def get_cardioblast_nuclei_data(
 88    path: Union[os.PathLike, str],
 89    split: Literal["train", "test"] = "train",
 90    download: bool = False,
 91) -> str:
 92    """Download the cardioblast nuclei dataset.
 93
 94    Args:
 95        path: Filepath to a folder where the downloaded data will be saved.
 96        split: The data split. Either 'train' or 'test'.
 97        download: Whether to download the data if it is not present.
 98
 99    Returns:
100        The filepath to the selected data split.
101    """
102    if split not in SAMPLES:
103        raise ValueError(f"'{split}' is not a valid split. Choose from {list(SAMPLES)}.")
104
105    split_dir = os.path.join(path, "cardioblast_nuclei", f"cardioblast_nuclei_{split}")
106    for sample in SAMPLES[split]:
107        sample_dir = os.path.join(split_dir, sample)
108        os.makedirs(sample_dir, exist_ok=True)
109
110        raw_path = os.path.join(sample_dir, f"{sample}.tif")
111        label_path = os.path.join(sample_dir, f"{sample}_mask.tif")
112        sample_url = f"{BASE_URL}/cardioblast_nuclei_{split}/{sample}"
113
114        util.download_source(raw_path, f"{sample_url}/{sample}.tif", download, checksum=RAW_CHECKSUMS[sample])
115        util.download_source(
116            label_path, f"{sample_url}/{sample}_mask.tif", download, checksum=LABEL_CHECKSUMS[sample]
117        )
118
119    return split_dir
120
121
122def get_cardioblast_nuclei_paths(
123    path: Union[os.PathLike, str],
124    split: Literal["train", "test"] = "train",
125    download: bool = False,
126) -> Tuple[List[str], List[str]]:
127    """Get paths to the cardioblast images and nucleus instance labels.
128
129    Args:
130        path: Filepath to a folder where the downloaded data will be saved.
131        split: The data split. Either 'train' or 'test'.
132        download: Whether to download the data if it is not present.
133
134    Returns:
135        The image paths and corresponding label paths.
136    """
137    split_dir = get_cardioblast_nuclei_data(path, split, download)
138    raw_paths = [os.path.join(split_dir, sample, f"{sample}.tif") for sample in SAMPLES[split]]
139    label_paths = [os.path.join(split_dir, sample, f"{sample}_mask.tif") for sample in SAMPLES[split]]
140
141    missing_paths = [path for path in raw_paths + label_paths if not os.path.exists(path)]
142    if missing_paths:
143        raise RuntimeError(f"Could not find {len(missing_paths)} cardioblast nuclei files for split '{split}'.")
144
145    return raw_paths, label_paths
146
147
148def get_cardioblast_nuclei_dataset(
149    path: Union[os.PathLike, str],
150    patch_shape: Tuple[int, int],
151    split: Literal["train", "test"] = "train",
152    offsets: Optional[List[List[int]]] = None,
153    boundaries: bool = False,
154    binary: bool = False,
155    download: bool = False,
156    **kwargs,
157) -> Dataset:
158    """Get the cardioblast nuclei dataset for instance segmentation.
159
160    Args:
161        path: Filepath to a folder where the downloaded data will be saved.
162        patch_shape: The 2D patch shape to use for training.
163        split: The data split. Either 'train' or 'test'.
164        offsets: Offset values for affinity computation used as target.
165        boundaries: Whether to compute boundaries as the target.
166        binary: Whether to use a binary segmentation target.
167        download: Whether to download the data if it is not present.
168        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
169
170    Returns:
171        The segmentation dataset.
172    """
173    if len(patch_shape) != 2:
174        raise ValueError(f"The cardioblast nuclei patch shape must be two-dimensional, got {patch_shape}.")
175
176    raw_paths, label_paths = get_cardioblast_nuclei_paths(path, split, download)
177    rois = [
178        (slice(0, ANNOTATED_FRAMES.get(sample)), slice(None), slice(None))
179        for sample in SAMPLES[split]
180    ]
181    kwargs.setdefault("rois", rois)
182    kwargs, _ = util.add_instance_label_transform(
183        kwargs, add_binary_target=True, offsets=offsets, boundaries=boundaries, binary=binary,
184    )
185    kwargs = util.ensure_transforms(ndim=2, **kwargs)
186
187    return torch_em.default_segmentation_dataset(
188        raw_paths=raw_paths,
189        raw_key=None,
190        label_paths=label_paths,
191        label_key=None,
192        patch_shape=(1,) + patch_shape,
193        is_seg_dataset=True,
194        ndim=2,
195        **kwargs,
196    )
197
198
199def get_cardioblast_nuclei_loader(
200    path: Union[os.PathLike, str],
201    batch_size: int,
202    patch_shape: Tuple[int, int],
203    split: Literal["train", "test"] = "train",
204    offsets: Optional[List[List[int]]] = None,
205    boundaries: bool = False,
206    binary: bool = False,
207    download: bool = False,
208    **kwargs,
209) -> DataLoader:
210    """Get the cardioblast nuclei dataloader for instance segmentation.
211
212    Args:
213        path: Filepath to a folder where the downloaded data will be saved.
214        batch_size: The batch size for training.
215        patch_shape: The 2D patch shape to use for training.
216        split: The data split. Either 'train' or 'test'.
217        offsets: Offset values for affinity computation used as target.
218        boundaries: Whether to compute boundaries as the target.
219        binary: Whether to use a binary segmentation target.
220        download: Whether to download the data if it is not present.
221        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or the PyTorch DataLoader.
222
223    Returns:
224        The DataLoader.
225    """
226    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
227    dataset = get_cardioblast_nuclei_dataset(
228        path=path,
229        patch_shape=patch_shape,
230        split=split,
231        offsets=offsets,
232        boundaries=boundaries,
233        binary=binary,
234        download=download,
235        **ds_kwargs,
236    )
237    return torch_em.get_data_loader(dataset, batch_size=batch_size, **loader_kwargs)