torch_em.data.datasets.light_microscopy.nuc_morph_timelapse

  1"""The NucMorph timelapse dataset contains 3D fluorescence microscopy timelapses of hiPSC colonies
  2with nuclear instance segmentation annotations.
  3
  4The dataset holds 14 colonies over six conditions. Each colony provides a raw timelapse with a
  5Lamin B1 EGFP channel and a brightfield channel, and a matching nuclear instance segmentation. The
  6annotations come from a Vision Transformer based segmentation model, and they are much cleaner than
  7the ones of the related `nuc_morph` dataset.
  8
  9NOTE: One timepoint holds about 190 MB per array, so the loader downloads a subset of the
 10timepoints. Use `stride` to set how many timepoints it skips, or pass `timepoints` to select them.
 11
 12NOTE: The raw level 0 and the segmentation level 1 share one grid. The segmentation level 0 is an
 13upsampled version with a different shape, so it does not match the raw data.
 14
 15NOTE: The index of the EGFP channel differs per colony, and the segmentation of most colonies stops
 16before the raw data ends. This module stores both facts in `COLONIES`, and it pairs the timepoints
 17without an offset, which was verified against the image data.
 18
 19The dataset is located at https://open.quiltdata.com/b/allencell/tree/aics/nuc-morph-dataset/ under
 20the Allen Institute for Cell Science Terms of Use.
 21This dataset is from the publication https://doi.org/10.1016/j.cels.2025.101265.
 22Please cite it if you use this dataset in your research.
 23"""
 24
 25import os
 26from glob import glob
 27from natsort import natsorted
 28from typing import List, Literal, Optional, Sequence, Tuple, Union
 29
 30import numpy as np
 31
 32from torch.utils.data import DataLoader, Dataset
 33
 34import torch_em
 35
 36from .. import util
 37
 38
 39S3_BASE = (
 40    "https://allencell.s3.amazonaws.com/aics/nuc-morph-dataset/hipsc_fov_nuclei_timelapse_dataset/"
 41    "hipsc_fov_nuclei_timelapse_data_used_for_analysis"
 42)
 43
 44# colony -> (condition, index of the EGFP channel in the raw data)
 45COLONIES = {
 46    "20200323_05_large": ("baseline_colonies", 0),
 47    "20200323_06_medium": ("baseline_colonies", 0),
 48    "20200323_09_small": ("baseline_colonies", 0),
 49    "20220411_03_control": ("dna_replication_inhibitor", 0),
 50    "20220411_05_aphidicolin": ("dna_replication_inhibitor", 0),
 51    "20230424_01_control": ("dna_replication_inhibitor", 1),
 52    "20230424_03_control": ("dna_replication_inhibitor", 1),
 53    "20230424_05_aphidicolin": ("dna_replication_inhibitor", 1),
 54    "20230720_01_control": ("feeding_control", 1),
 55    "20230720_04_pre-starved": ("feeding_control", 1),
 56    "20230720_07_re-fed": ("feeding_control", 1),
 57    "20220901_01": ("fixed_control", 1),
 58    "20230417_01_control": ("nuclear_import_inhibitor", 1),
 59    "20230417_07_importazole": ("nuclear_import_inhibitor", 1),
 60}
 61
 62CONDITIONS = tuple(dict.fromkeys(condition for condition, _ in COLONIES.values()))
 63
 64CHANNELS = ("egfp", "brightfield", "both")
 65
 66
 67def _open_array(url: str):
 68    """Open a remote zarr array over http."""
 69    import zarr
 70
 71    try:
 72        return zarr.open(url, mode="r")
 73    except Exception:
 74        from zarr.storage import FsspecStore
 75        return zarr.open(store=FsspecStore.from_url(url), mode="r")
 76
 77
 78def _get_colonies(condition: Optional[str], colony: Optional[Union[str, Sequence[str]]]) -> List[str]:
 79    """Resolve the requested colonies."""
 80    if colony is not None:
 81        colonies = [colony] if isinstance(colony, str) else list(colony)
 82        for name in colonies:
 83            if name not in COLONIES:
 84                raise ValueError(f"'{name}' is not a valid colony. Choose from {list(COLONIES)}.")
 85        return colonies
 86
 87    if condition is None:
 88        return list(COLONIES)
 89
 90    if condition not in CONDITIONS:
 91        raise ValueError(f"'{condition}' is not a valid condition. Choose from {list(CONDITIONS)}.")
 92    return [name for name, (this_condition, _) in COLONIES.items() if this_condition == condition]
 93
 94
 95def _download_colony(
 96    path: str, colony: str, timepoints: Optional[Sequence[int]], stride: int, channel: str, download: bool,
 97) -> str:
 98    """Download the selected timepoints of one colony and store them as h5 files."""
 99    import h5py
100    from tqdm import tqdm
101
102    condition, egfp_channel = COLONIES[colony]
103    colony_dir = os.path.join(path, colony)
104    os.makedirs(colony_dir, exist_ok=True)
105
106    base = f"{S3_BASE}/{condition}_fov_timelapse_dataset/{colony}"
107    raw_array = _open_array(f"{base}/raw.ome.zarr/0")
108    seg_array = _open_array(f"{base}/seg.ome.zarr/1")
109
110    if raw_array.shape[-3:] != seg_array.shape[-3:]:
111        raise RuntimeError(
112            f"The raw and the segmentation grid of '{colony}' differ, "
113            f"{raw_array.shape[-3:]} against {seg_array.shape[-3:]}."
114        )
115
116    # The segmentation stops before the raw data ends, so it limits the valid timepoints.
117    n_timepoints = seg_array.shape[0]
118    if timepoints is None:
119        timepoints = range(0, n_timepoints, stride)
120    selected = [int(t) for t in timepoints]
121    for timepoint in selected:
122        if not 0 <= timepoint < n_timepoints:
123            raise ValueError(f"The timepoint {timepoint} is outside the segmented range of '{colony}', "
124                             f"which holds {n_timepoints} timepoints.")
125
126    if channel == "egfp":
127        channel_ids = [egfp_channel]
128    elif channel == "brightfield":
129        channel_ids = [1 - egfp_channel]
130    else:
131        channel_ids = [egfp_channel, 1 - egfp_channel]
132
133    for timepoint in tqdm(selected, desc=f"Download '{colony}'"):
134        output_path = os.path.join(colony_dir, f"t{timepoint:04d}.h5")
135        if os.path.exists(output_path):
136            continue
137
138        if not download:
139            raise RuntimeError(f"Cannot find the data at {output_path}, but download was set to False.")
140
141        raw = np.stack([np.asarray(raw_array[timepoint, c]) for c in channel_ids])
142        labels = np.asarray(seg_array[timepoint, 0])
143        if raw.shape[0] == 1:
144            raw = raw[0]
145
146        with h5py.File(output_path, "w") as f:
147            f.create_dataset("raw", data=raw, compression="gzip")
148            f.create_dataset("labels", data=labels, compression="gzip")
149
150    return colony_dir
151
152
153def get_nuc_morph_timelapse_data(
154    path: Union[os.PathLike, str],
155    condition: Optional[str] = "baseline_colonies",
156    colony: Optional[Union[str, Sequence[str]]] = None,
157    timepoints: Optional[Sequence[int]] = None,
158    stride: int = 50,
159    channel: Literal["egfp", "brightfield", "both"] = "egfp",
160    download: bool = False,
161) -> List[str]:
162    """Download the NucMorph timelapse dataset.
163
164    Args:
165        path: Filepath to a folder where the downloaded data will be saved.
166        condition: The experimental condition. Ignored when you pass `colony`.
167        colony: The colony or colonies to use. Overrides `condition`.
168        timepoints: The timepoints to download. Overrides `stride`.
169        stride: The number of timepoints to skip between two downloads.
170        channel: The raw channel. Either 'egfp', 'brightfield' or 'both'.
171        download: Whether to download the data if it is not present.
172
173    Returns:
174        List of the folders that hold the data of the requested colonies.
175    """
176    if channel not in CHANNELS:
177        raise ValueError(f"'{channel}' is not a valid channel. Choose from {list(CHANNELS)}.")
178    if stride < 1:
179        raise ValueError(f"The stride must be at least one, got {stride}.")
180
181    colonies = _get_colonies(condition, colony)
182    os.makedirs(path, exist_ok=True)
183    return [_download_colony(path, name, timepoints, stride, channel, download) for name in colonies]
184
185
186def get_nuc_morph_timelapse_paths(
187    path: Union[os.PathLike, str],
188    condition: Optional[str] = "baseline_colonies",
189    colony: Optional[Union[str, Sequence[str]]] = None,
190    timepoints: Optional[Sequence[int]] = None,
191    stride: int = 50,
192    channel: Literal["egfp", "brightfield", "both"] = "egfp",
193    download: bool = False,
194) -> List[str]:
195    """Get paths to the NucMorph timelapse data.
196
197    Args:
198        path: Filepath to a folder where the downloaded data will be saved.
199        condition: The experimental condition. Ignored when you pass `colony`.
200        colony: The colony or colonies to use. Overrides `condition`.
201        timepoints: The timepoints to download. Overrides `stride`.
202        stride: The number of timepoints to skip between two downloads.
203        channel: The raw channel. Either 'egfp', 'brightfield' or 'both'.
204        download: Whether to download the data if it is not present.
205
206    Returns:
207        List of filepaths for the h5 data.
208    """
209    colony_dirs = get_nuc_morph_timelapse_data(path, condition, colony, timepoints, stride, channel, download)
210
211    volume_paths = []
212    for colony_dir in colony_dirs:
213        volume_paths.extend(natsorted(glob(os.path.join(colony_dir, "*.h5"))))
214
215    if not volume_paths:
216        raise RuntimeError(f"Could not find any NucMorph timelapse data in {path}.")
217
218    return volume_paths
219
220
221def get_nuc_morph_timelapse_dataset(
222    path: Union[os.PathLike, str],
223    patch_shape: Tuple[int, int, int],
224    condition: Optional[str] = "baseline_colonies",
225    colony: Optional[Union[str, Sequence[str]]] = None,
226    timepoints: Optional[Sequence[int]] = None,
227    stride: int = 50,
228    channel: Literal["egfp", "brightfield", "both"] = "egfp",
229    offsets: Optional[List[List[int]]] = None,
230    boundaries: bool = False,
231    binary: bool = False,
232    download: bool = False,
233    **kwargs,
234) -> Dataset:
235    """Get the NucMorph timelapse dataset for nucleus segmentation.
236
237    Args:
238        path: Filepath to a folder where the downloaded data will be saved.
239        patch_shape: The 3D patch shape to use for training.
240        condition: The experimental condition. Ignored when you pass `colony`.
241        colony: The colony or colonies to use. Overrides `condition`.
242        timepoints: The timepoints to download. Overrides `stride`.
243        stride: The number of timepoints to skip between two downloads.
244        channel: The raw channel. Either 'egfp', 'brightfield' or 'both'.
245        offsets: Offset values for affinity computation used as target.
246        boundaries: Whether to compute boundaries as the target.
247        binary: Whether to use a binary segmentation target.
248        download: Whether to download the data if it is not present.
249        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
250
251    Returns:
252        The segmentation dataset.
253    """
254    if len(patch_shape) != 3:
255        raise ValueError(f"The NucMorph timelapse patch shape must be three-dimensional, got {patch_shape}.")
256
257    volume_paths = get_nuc_morph_timelapse_paths(
258        path, condition, colony, timepoints, stride, channel, download
259    )
260
261    kwargs, _ = util.add_instance_label_transform(
262        kwargs, add_binary_target=True, offsets=offsets, boundaries=boundaries, binary=binary,
263    )
264    kwargs = util.ensure_transforms(ndim=3, **kwargs)
265
266    return torch_em.default_segmentation_dataset(
267        raw_paths=volume_paths,
268        raw_key="raw",
269        label_paths=volume_paths,
270        label_key="labels",
271        patch_shape=patch_shape,
272        ndim=3,
273        with_channels=channel == "both",
274        **kwargs,
275    )
276
277
278def get_nuc_morph_timelapse_loader(
279    path: Union[os.PathLike, str],
280    batch_size: int,
281    patch_shape: Tuple[int, int, int],
282    condition: Optional[str] = "baseline_colonies",
283    colony: Optional[Union[str, Sequence[str]]] = None,
284    timepoints: Optional[Sequence[int]] = None,
285    stride: int = 50,
286    channel: Literal["egfp", "brightfield", "both"] = "egfp",
287    offsets: Optional[List[List[int]]] = None,
288    boundaries: bool = False,
289    binary: bool = False,
290    download: bool = False,
291    **kwargs,
292) -> DataLoader:
293    """Get the NucMorph timelapse dataloader for nucleus segmentation.
294
295    Args:
296        path: Filepath to a folder where the downloaded data will be saved.
297        batch_size: The batch size for training.
298        patch_shape: The 3D patch shape to use for training.
299        condition: The experimental condition. Ignored when you pass `colony`.
300        colony: The colony or colonies to use. Overrides `condition`.
301        timepoints: The timepoints to download. Overrides `stride`.
302        stride: The number of timepoints to skip between two downloads.
303        channel: The raw channel. Either 'egfp', 'brightfield' or 'both'.
304        offsets: Offset values for affinity computation used as target.
305        boundaries: Whether to compute boundaries as the target.
306        binary: Whether to use a binary segmentation target.
307        download: Whether to download the data if it is not present.
308        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or the PyTorch DataLoader.
309
310    Returns:
311        The DataLoader.
312    """
313    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
314    dataset = get_nuc_morph_timelapse_dataset(
315        path=path,
316        patch_shape=patch_shape,
317        condition=condition,
318        colony=colony,
319        timepoints=timepoints,
320        stride=stride,
321        channel=channel,
322        offsets=offsets,
323        boundaries=boundaries,
324        binary=binary,
325        download=download,
326        **ds_kwargs,
327    )
328    return torch_em.get_data_loader(dataset, batch_size=batch_size, **loader_kwargs)