torch_em.data.datasets.light_microscopy.nucverse3d

  1"""The NucVerse3D dataset contains annotations for 3d nucleus segmentation in
  2two-photon and confocal microscopy volumes.
  3
  4The dataset holds three collections that the authors annotated by hand:
  5- 'liver': two-photon volumes of adult mouse liver, with mono- and binucleated hepatocytes,
  6- 'liver_hcc': two-photon volumes of mouse hepatocellular carcinoma,
  7- 'drosophila_glia': confocal volumes of Drosophila melanogaster brain, stained for glial nuclei.
  8
  9Together they hold 26 volumes with 6226 nuclei. Every volume becomes one h5 file that stores the
 10raw data, the instance labels and the voxel size.
 11
 12NOTE: The three collections name their masks in three ways. The liver collection appends '_gt' or
 13'_gt_uint16', the carcinoma collection appends '_labels', and the fly collection replaces a
 14'Denoised-' prefix with a '-label' suffix. The loader therefore pairs an image with a mask by
 15position after sorting, and it checks that the shape of a pair agrees.
 16
 17NOTE: The masks come as int32, uint32 and uint16, even inside one split. The loader writes them all
 18as uint32.
 19
 20NOTE: The readme of the repository says that the carcinoma collection serves for testing alone, but
 21the archive splits it into 8 training volumes and 4 test volumes. This loader follows the archive.
 22
 23NOTE: The tiff files carry no usable voxel size. The carcinoma files report a resolution of one with
 24the unit 'none', and the fly files report no resolution at all, so the voxel size of this module
 25comes from the readme of the repository.
 26
 27The dataset is located at https://doi.org/10.5281/zenodo.18517324 under the CC BY 4.0 license.
 28This dataset is from the publication https://doi.org/10.1038/s41598-026-51994-x.
 29Please cite it if you use this dataset in your research.
 30"""
 31
 32import os
 33from glob import glob
 34from natsort import natsorted
 35from typing import List, Literal, Optional, Sequence, Tuple, Union
 36
 37from torch.utils.data import DataLoader, Dataset
 38
 39import torch_em
 40
 41from .. import util
 42
 43
 44URL = "https://zenodo.org/records/18517324/files/raw.zip?download=1"
 45CHECKSUM = "ba4be8b887fe839bb4b16f50e87b93292f95b959e9dc2d161e234521cc9a6d19"
 46
 47# The name of a collection in the archive, and what the readme says about it.
 48DATASETS = {
 49    "liver": {
 50        "folder": "2_livernuclei",
 51        "modality": "two-photon microscopy",
 52        "tissue": "adult mouse liver",
 53        # The readme gives the voxel size as x, y, z in micrometer. The arrays are z, y, x.
 54        "resolution": (0.3, 0.3, 0.3),
 55    },
 56    "liver_hcc": {
 57        "folder": "7_liver_hcc_dataset",
 58        "modality": "two-photon microscopy",
 59        "tissue": "mouse liver hepatocellular carcinoma",
 60        "resolution": (0.3, 0.3, 0.3),
 61    },
 62    "drosophila_glia": {
 63        "folder": "6_Drosophila_denoised",
 64        "modality": "confocal microscopy",
 65        "tissue": "Drosophila melanogaster brain glia",
 66        "resolution": (0.3, 0.15, 0.15),
 67    },
 68}
 69
 70SPLITS = ("train", "test")
 71
 72
 73def _pair_split(data_dir: str, folder: str, split: str) -> List[Tuple[str, str]]:
 74    """Pair every image of a split with its mask, by position after sorting."""
 75    image_paths = natsorted(glob(os.path.join(data_dir, folder, split, "images", "*.tif")))
 76    mask_paths = natsorted(glob(os.path.join(data_dir, folder, split, "masks", "*.tif")))
 77    if len(image_paths) != len(mask_paths):
 78        raise RuntimeError(
 79            f"The split '{folder}/{split}' holds {len(image_paths)} images but {len(mask_paths)} masks."
 80        )
 81    return list(zip(image_paths, mask_paths))
 82
 83
 84def _create_h5(data_dir: str, name: str, split: str) -> str:
 85    """Write one h5 file per volume, with the raw data, the labels and the voxel size."""
 86    import h5py
 87    import tifffile
 88    from tqdm import tqdm
 89
 90    info = DATASETS[name]
 91    output_dir = os.path.join(data_dir, "preprocessed", name, split)
 92    os.makedirs(output_dir, exist_ok=True)
 93
 94    pairs = _pair_split(data_dir, info["folder"], split)
 95    for image_path, mask_path in tqdm(pairs, desc=f"Preprocess '{name}/{split}'"):
 96        stem = os.path.splitext(os.path.basename(image_path))[0]
 97        output_path = os.path.join(output_dir, f"{stem}.h5")
 98        if os.path.exists(output_path):
 99            continue
100
101        raw = tifffile.imread(image_path)
102        labels = tifffile.imread(mask_path)
103        if raw.shape != labels.shape:
104            raise RuntimeError(
105                f"The image {os.path.basename(image_path)} has the shape {raw.shape}, "
106                f"but its mask {os.path.basename(mask_path)} has the shape {labels.shape}."
107            )
108
109        temporary_path = f"{output_path}.tmp"
110        with h5py.File(temporary_path, "w") as f:
111            f.attrs["dataset"] = name
112            f.attrs["modality"] = info["modality"]
113            f.attrs["tissue"] = info["tissue"]
114            f.attrs["split"] = split
115            # The voxel size in micrometer, in the order of the axes of the arrays.
116            f.attrs["resolution"] = info["resolution"]
117            f.attrs["axes"] = "zyx"
118            f.attrs["image_file"] = os.path.basename(image_path)
119            f.attrs["label_file"] = os.path.basename(mask_path)
120
121            raw_dataset = f.create_dataset("raw", data=raw, compression="gzip")
122            raw_dataset.attrs["resolution"] = info["resolution"]
123            # The masks come in three integer types, so they all become uint32 here.
124            label_dataset = f.create_dataset("labels", data=labels.astype("uint32"), compression="gzip")
125            label_dataset.attrs["resolution"] = info["resolution"]
126        os.replace(temporary_path, output_path)
127
128    return output_dir
129
130
131def get_nucverse3d_data(path: Union[os.PathLike, str], download: bool = False) -> str:
132    """Download the NucVerse3D dataset.
133
134    The repository also holds trained models and preprocessed patches, which take 31 GB together.
135    This loader reads the raw archive alone, which takes 0.68 GB.
136
137    Args:
138        path: Filepath to a folder where the downloaded data will be saved.
139        download: Whether to download the data if it is not present.
140
141    Returns:
142        The filepath to the extracted data.
143    """
144    data_dir = os.path.join(path, "raw")
145    if os.path.exists(data_dir):
146        return data_dir
147
148    os.makedirs(path, exist_ok=True)
149    zip_path = os.path.join(path, "raw.zip")
150    util.download_source(zip_path, URL, download, CHECKSUM)
151    util.unzip(zip_path=zip_path, dst=path)
152
153    return data_dir
154
155
156def get_nucverse3d_paths(
157    path: Union[os.PathLike, str],
158    dataset: Optional[Union[str, Sequence[str]]] = None,
159    split: Literal["train", "test"] = "train",
160    download: bool = False,
161) -> List[str]:
162    """Get paths to the NucVerse3D data.
163
164    Args:
165        path: Filepath to a folder where the downloaded data will be saved.
166        dataset: The collection or collections to use, out of 'liver', 'liver_hcc' and
167            'drosophila_glia'. Defaults to all of them.
168        split: The data split. Either 'train' or 'test'.
169        download: Whether to download the data if it is not present.
170
171    Returns:
172        List of filepaths for the h5 data.
173    """
174    if split not in SPLITS:
175        raise ValueError(f"'{split}' is not a valid split. Choose from {list(SPLITS)}.")
176
177    if dataset is None:
178        names = list(DATASETS)
179    else:
180        names = [dataset] if isinstance(dataset, str) else list(dataset)
181        for name in names:
182            if name not in DATASETS:
183                raise ValueError(f"'{name}' is not a valid dataset. Choose from {list(DATASETS)}.")
184
185    data_dir = get_nucverse3d_data(path, download)
186
187    volume_paths = []
188    for name in names:
189        output_dir = _create_h5(data_dir, name, split)
190        volume_paths.extend(natsorted(glob(os.path.join(output_dir, "*.h5"))))
191
192    if not volume_paths:
193        raise RuntimeError(f"Could not find any NucVerse3D data in {data_dir}.")
194
195    return volume_paths
196
197
198def get_nucverse3d_dataset(
199    path: Union[os.PathLike, str],
200    patch_shape: Tuple[int, int, int],
201    dataset: Optional[Union[str, Sequence[str]]] = None,
202    split: Literal["train", "test"] = "train",
203    offsets: Optional[List[List[int]]] = None,
204    boundaries: bool = False,
205    binary: bool = False,
206    download: bool = False,
207    **kwargs,
208) -> Dataset:
209    """Get the NucVerse3D dataset for 3d nucleus segmentation.
210
211    Args:
212        path: Filepath to a folder where the downloaded data will be saved.
213        patch_shape: The 3D patch shape to use for training.
214        dataset: The collection or collections to use, out of 'liver', 'liver_hcc' and
215            'drosophila_glia'. Defaults to all of them.
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`.
222
223    Returns:
224        The segmentation dataset.
225    """
226    if len(patch_shape) != 3:
227        raise ValueError(f"The NucVerse3D patch shape must be three-dimensional, got {patch_shape}.")
228
229    volume_paths = get_nucverse3d_paths(path, dataset, split, download)
230
231    kwargs, _ = util.add_instance_label_transform(
232        kwargs, add_binary_target=True, offsets=offsets, boundaries=boundaries, binary=binary,
233    )
234    kwargs = util.ensure_transforms(ndim=3, **kwargs)
235
236    return torch_em.default_segmentation_dataset(
237        raw_paths=volume_paths,
238        raw_key="raw",
239        label_paths=volume_paths,
240        label_key="labels",
241        patch_shape=patch_shape,
242        ndim=3,
243        **kwargs,
244    )
245
246
247def get_nucverse3d_loader(
248    path: Union[os.PathLike, str],
249    batch_size: int,
250    patch_shape: Tuple[int, int, int],
251    dataset: Optional[Union[str, Sequence[str]]] = None,
252    split: Literal["train", "test"] = "train",
253    offsets: Optional[List[List[int]]] = None,
254    boundaries: bool = False,
255    binary: bool = False,
256    download: bool = False,
257    **kwargs,
258) -> DataLoader:
259    """Get the NucVerse3D dataloader for 3d nucleus segmentation.
260
261    Args:
262        path: Filepath to a folder where the downloaded data will be saved.
263        batch_size: The batch size for training.
264        patch_shape: The 3D patch shape to use for training.
265        dataset: The collection or collections to use, out of 'liver', 'liver_hcc' and
266            'drosophila_glia'. Defaults to all of them.
267        split: The data split. Either 'train' or 'test'.
268        offsets: Offset values for affinity computation used as target.
269        boundaries: Whether to compute boundaries as the target.
270        binary: Whether to use a binary segmentation target.
271        download: Whether to download the data if it is not present.
272        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or the PyTorch DataLoader.
273
274    Returns:
275        The DataLoader.
276    """
277    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
278    dataset_ = get_nucverse3d_dataset(
279        path=path,
280        patch_shape=patch_shape,
281        dataset=dataset,
282        split=split,
283        offsets=offsets,
284        boundaries=boundaries,
285        binary=binary,
286        download=download,
287        **ds_kwargs,
288    )
289    return torch_em.get_data_loader(dataset_, batch_size=batch_size, **loader_kwargs)