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, Sequence
 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    volumes: Optional[Sequence[str]] = None,
106    download: bool = False,
107) -> List[str]:
108    """Get paths to the Wing Disc data.
109
110    Args:
111        path: Filepath to a folder where the downloaded data will be saved.
112        volumes: The volume names to restrict to, see VOLUMES. By default all four volumes are used.
113        download: Whether to download the data if it is not present.
114
115    Returns:
116        List of filepaths for the stored data.
117    """
118    data_dir = get_wing_disc_data(path, download)
119    data_paths = natsorted(glob(os.path.join(data_dir, "*.h5")))
120    if volumes is not None:
121        data_paths = [p for p in data_paths if os.path.splitext(os.path.basename(p))[0] in volumes]
122    assert len(data_paths) > 0
123    return data_paths
124
125
126def get_wing_disc_dataset(
127    path: Union[os.PathLike, str],
128    patch_shape: Tuple[int, int, int],
129    offsets: Optional[List[List[int]]] = None,
130    boundaries: bool = False,
131    binary: bool = False,
132    volumes: Optional[Sequence[str]] = None,
133    download: bool = False,
134    **kwargs
135) -> Dataset:
136    """Get the Wing Disc dataset for 3D cell segmentation in Drosophila wing discs.
137
138    Args:
139        path: Filepath to a folder where the downloaded data will be saved.
140        patch_shape: The patch shape to use for training.
141        offsets: Offset values for affinity computation used as target.
142        boundaries: Whether to compute boundaries as the target.
143        binary: Whether to use a binary segmentation target.
144        volumes: The volume names to restrict to, see VOLUMES. By default all four volumes are used.
145        download: Whether to download the data if it is not present.
146        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
147
148    Returns:
149        The segmentation dataset.
150    """
151    data_paths = get_wing_disc_paths(path, volumes, download)
152
153    kwargs = util.ensure_transforms(ndim=3, **kwargs)
154    kwargs, _ = util.add_instance_label_transform(
155        kwargs, add_binary_target=True, offsets=offsets, boundaries=boundaries, binary=binary
156    )
157
158    return torch_em.default_segmentation_dataset(
159        raw_paths=data_paths,
160        raw_key="raw",
161        label_paths=data_paths,
162        label_key="labels",
163        patch_shape=patch_shape,
164        ndim=3,
165        **kwargs
166    )
167
168
169def get_wing_disc_loader(
170    path: Union[os.PathLike, str],
171    batch_size: int,
172    patch_shape: Tuple[int, int, int],
173    offsets: Optional[List[List[int]]] = None,
174    boundaries: bool = False,
175    binary: bool = False,
176    volumes: Optional[Sequence[str]] = None,
177    download: bool = False,
178    **kwargs
179) -> DataLoader:
180    """Get the Wing Disc dataloader for 3D cell segmentation in Drosophila wing discs.
181
182    Args:
183        path: Filepath to a folder where the downloaded data will be saved.
184        batch_size: The batch size for training.
185        patch_shape: The patch shape to use for training.
186        offsets: Offset values for affinity computation used as target.
187        boundaries: Whether to compute boundaries as the target.
188        binary: Whether to use a binary segmentation target.
189        volumes: The volume names to restrict to, see VOLUMES. By default all four volumes are used.
190        download: Whether to download the data if it is not present.
191        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
192
193    Returns:
194        The DataLoader.
195    """
196    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
197    dataset = get_wing_disc_dataset(
198        path=path,
199        patch_shape=patch_shape,
200        offsets=offsets,
201        boundaries=boundaries,
202        binary=binary,
203        volumes=volumes,
204        download=download,
205        **ds_kwargs,
206    )
207    return torch_em.get_data_loader(dataset=dataset, batch_size=batch_size, **loader_kwargs)