torch_em.data.datasets.light_microscopy.wing_disc

  1"""The Wing Disc dataset contains annotations for 3D cell instance segmentation
  2in confocal microscopy images of Drosophila wing discs.
  3
  4The dataset is located at https://www.ebi.ac.uk/biostudies/BioImages/studies/S-BIAD843.
  5This dataset is from the publication https://www.nature.com/articles/s44303-025-00099-7.
  6Please cite it if you use this dataset in your research.
  7"""
  8
  9import os
 10from glob import glob
 11from natsort import natsorted
 12from typing import Union, Tuple, Optional, List
 13
 14import numpy as np
 15
 16from torch.utils.data import Dataset, DataLoader
 17
 18import torch_em
 19
 20from .. import util
 21
 22
 23BASE_URL = "https://ftp.ebi.ac.uk/biostudies/fire/S-BIAD/843/S-BIAD843/Files"
 24
 25VOLUMES = {
 26    "WD1_15-02_WT_confocalonly": "confocal",
 27    "WD2.1_21-02_WT_confocalonly": "confocal",
 28    "WD1.1_17-03_WT_MP": "multiphoton",
 29    "WD3.2_21-03_WT_MP": "multiphoton",
 30}
 31
 32
 33def _preprocess_volumes(path, data_dir):
 34    """Convert OME-Zarr volumes to HDF5 files with raw and labels datasets."""
 35    import h5py
 36    import zarr
 37
 38    os.makedirs(data_dir, exist_ok=True)
 39
 40    zarr_dir = os.path.join(path, "zarr")
 41
 42    for name in VOLUMES:
 43        h5_path = os.path.join(data_dir, f"{name}.h5")
 44        if os.path.exists(h5_path):
 45            continue
 46
 47        # Read raw volume: shape (1, 1, Z, Y, X) and squeeze to (Z, Y, X).
 48        raw_zarr = os.path.join(zarr_dir, f"{name}.zarr", "0", "0")
 49        raw = np.array(zarr.open(store=zarr.storage.LocalStore(raw_zarr)))
 50        raw = raw.squeeze()
 51
 52        # Read segmentation: shape (Z, 1, 1, Y, X) and squeeze to (Z, Y, X).
 53        seg_zarr = os.path.join(zarr_dir, f"{name}_segmented.zarr", "0", "0")
 54        seg = np.array(zarr.open(store=zarr.storage.LocalStore(seg_zarr)))
 55        seg = seg.squeeze().astype("uint32")
 56
 57        assert raw.shape == seg.shape, f"Shape mismatch for {name}: raw={raw.shape}, seg={seg.shape}"
 58
 59        with h5py.File(h5_path, "w") as f:
 60            f.create_dataset("raw", data=raw, compression="gzip")
 61            f.create_dataset("labels", data=seg, compression="gzip")
 62
 63
 64def get_wing_disc_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 65    """Download the Wing Disc dataset.
 66
 67    Args:
 68        path: Filepath to a folder where the downloaded data will be saved.
 69        download: Whether to download the data if it is not present.
 70
 71    Returns:
 72        The filepath to the preprocessed data directory.
 73    """
 74    data_dir = os.path.join(path, "data")
 75    if os.path.exists(data_dir) and len(glob(os.path.join(data_dir, "*.h5"))) == len(VOLUMES):
 76        return data_dir
 77
 78    zarr_dir = os.path.join(path, "zarr")
 79    os.makedirs(zarr_dir, exist_ok=True)
 80
 81    for name in VOLUMES:
 82        zarr_path = os.path.join(zarr_dir, f"{name}.zarr")
 83        if not os.path.exists(zarr_path):
 84            zip_fname = f"{name}.ome.zarr.zip"
 85            zip_path = os.path.join(path, zip_fname)
 86            url = f"{BASE_URL}/{zip_fname}"
 87            util.download_source(path=zip_path, url=url, download=download, checksum=None)
 88            util.unzip(zip_path=zip_path, dst=zarr_dir)
 89
 90        seg_zarr_path = os.path.join(zarr_dir, f"{name}_segmented.zarr")
 91        if not os.path.exists(seg_zarr_path):
 92            seg_zip_fname = f"{name}_segmented.ome.zarr.zip"
 93            seg_zip_path = os.path.join(path, seg_zip_fname)
 94            seg_url = f"{BASE_URL}/{seg_zip_fname}"
 95            util.download_source(path=seg_zip_path, url=seg_url, download=download, checksum=None)
 96            util.unzip(zip_path=seg_zip_path, dst=zarr_dir)
 97
 98    _preprocess_volumes(path, data_dir)
 99
100    return data_dir
101
102
103def get_wing_disc_paths(
104    path: Union[os.PathLike, str],
105    download: bool = False,
106) -> List[str]:
107    """Get paths to the Wing Disc data.
108
109    Args:
110        path: Filepath to a folder where the downloaded data will be saved.
111        download: Whether to download the data if it is not present.
112
113    Returns:
114        List of filepaths for the stored data.
115    """
116    data_dir = get_wing_disc_data(path, download)
117    data_paths = natsorted(glob(os.path.join(data_dir, "*.h5")))
118    assert len(data_paths) > 0
119    return data_paths
120
121
122def get_wing_disc_dataset(
123    path: Union[os.PathLike, str],
124    patch_shape: Tuple[int, int, int],
125    offsets: Optional[List[List[int]]] = None,
126    boundaries: bool = False,
127    binary: bool = False,
128    download: bool = False,
129    **kwargs
130) -> Dataset:
131    """Get the Wing Disc dataset for 3D cell segmentation in Drosophila wing discs.
132
133    Args:
134        path: Filepath to a folder where the downloaded data will be saved.
135        patch_shape: The patch shape to use for training.
136        offsets: Offset values for affinity computation used as target.
137        boundaries: Whether to compute boundaries as the target.
138        binary: Whether to use a binary segmentation target.
139        download: Whether to download the data if it is not present.
140        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
141
142    Returns:
143        The segmentation dataset.
144    """
145    data_paths = get_wing_disc_paths(path, download)
146
147    kwargs = util.ensure_transforms(ndim=3, **kwargs)
148    kwargs, _ = util.add_instance_label_transform(
149        kwargs, add_binary_target=True, offsets=offsets, boundaries=boundaries, binary=binary
150    )
151
152    return torch_em.default_segmentation_dataset(
153        raw_paths=data_paths,
154        raw_key="raw",
155        label_paths=data_paths,
156        label_key="labels",
157        patch_shape=patch_shape,
158        ndim=3,
159        **kwargs
160    )
161
162
163def get_wing_disc_loader(
164    path: Union[os.PathLike, str],
165    batch_size: int,
166    patch_shape: Tuple[int, int, int],
167    offsets: Optional[List[List[int]]] = None,
168    boundaries: bool = False,
169    binary: bool = False,
170    download: bool = False,
171    **kwargs
172) -> DataLoader:
173    """Get the Wing Disc dataloader for 3D cell segmentation in Drosophila wing discs.
174
175    Args:
176        path: Filepath to a folder where the downloaded data will be saved.
177        batch_size: The batch size for training.
178        patch_shape: The patch shape to use for training.
179        offsets: Offset values for affinity computation used as target.
180        boundaries: Whether to compute boundaries as the target.
181        binary: Whether to use a binary segmentation target.
182        download: Whether to download the data if it is not present.
183        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
184
185    Returns:
186        The DataLoader.
187    """
188    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
189    dataset = get_wing_disc_dataset(
190        path=path,
191        patch_shape=patch_shape,
192        offsets=offsets,
193        boundaries=boundaries,
194        binary=binary,
195        download=download,
196        **ds_kwargs,
197    )
198    return torch_em.get_data_loader(dataset=dataset, batch_size=batch_size, **loader_kwargs)