torch_em.data.datasets.light_microscopy.selma3d

  1"""The SELMA3D dataset contains annotated 3D light-sheet microscopy images of cleared brain tissue.
  2
  3This loader provides the cell nucleus subset used by the model-ranking benchmark. It contains twelve
  4fluorescence volumes with binary nucleus annotations. The train, validation and test split follows the
  5benchmark split from https://github.com/kreshuklab/model_ranking.
  6
  7The dataset is located at https://doi.org/10.6019/S-BIAD1196 and is available under the CC BY 4.0 license.
  8It is from the publication https://doi.org/10.48550/arXiv.2501.03880.
  9Please cite the dataset and publication if you use this dataset in your research.
 10"""
 11
 12import os
 13from typing import List, Literal, Tuple, Union
 14
 15from torch.utils.data import DataLoader, Dataset
 16
 17import torch_em
 18
 19from .. import util
 20
 21
 22BASE_URL = "https://www.ebi.ac.uk/biostudies/files/S-BIAD1196"
 23DATA_ROOT = os.path.join("SELMA3D_training_annotated", "shannel_cells")
 24URL_ROOT = f"{BASE_URL}/SELMA3D_training_annotated/shannel_cells"
 25
 26SPLITS = {
 27    "train": tuple(f"patchvolume_{sample_id:03d}" for sample_id in range(8)),
 28    "val": ("patchvolume_008",),
 29    "test": tuple(f"patchvolume_{sample_id:03d}" for sample_id in range(9, 12)),
 30}
 31
 32RAW_CHECKSUMS = {
 33    "patchvolume_000": "25e1e351872db53a22f2ff9196892e4cefcc44c3d0be98186c13d3c77f8e0397",
 34    "patchvolume_001": "d12e43353982148b6726110ddef6331c0f057b964924accbd492c03ae2af09b7",
 35    "patchvolume_002": "74915c737bb470052808e3cdbeb6303fb4c64e81f0a0b0e375057116a9b10590",
 36    "patchvolume_003": "631a43a1f511f17ae2a2c92f2fb7841fd6f7eab683fcaf0ddaf530b7eb9a9501",
 37    "patchvolume_004": "f2043a0db9247d1a24349c499712223fc04050c54bbc2d09e4b854905faa5510",
 38    "patchvolume_005": "297e6e93418497a7f6af4814f08cd03741a21723a0e06ed7003db60ca59d8b9f",
 39    "patchvolume_006": "b2ec6c7d6d3a6c4d6fc99237a04c2ceb392aaed7505e33387bb10f11f1dc031a",
 40    "patchvolume_007": "f5e6219cc27a07772d696eb8f2c235b8f5aea96a1e6d1874a1e1254b38c90017",
 41    "patchvolume_008": "09bf8dac83abfeb575f76feaf4fb252258e78d7631b4c9f51a55460cfdaf9529",
 42    "patchvolume_009": "7f28921b85cb7a69c3937c3acfe1a4c1fa4091bd8d1d80a0c223e0f1ab82240c",
 43    "patchvolume_010": "acad1cdb850705d8263482c1a0c196f4ae801ddc6ef1a64e8f1b2132a93da486",
 44    "patchvolume_011": "acbb1c96da13958c892c9351aa0520988263640e3d247e961d93256bc2970f82",
 45}
 46
 47LABEL_CHECKSUMS = {
 48    "patchvolume_000": "ad59e596aba5bb3562360fa824f9f23474facf2d80242c79edcc5cd8b71f5491",
 49    "patchvolume_001": "9669072ce2a4cad1a96b5393afc0f10c6f146f3452ebfe6c6e7adbe941ca4eaa",
 50    "patchvolume_002": "3cebf4f81169b3349971b06a673d2d948bc3dfc929379c9d6664cb35553d8a98",
 51    "patchvolume_003": "648b71876b83c2d65c8cc5f6ecbaf2176e39b400483ebe233d7d17c52d95fa53",
 52    "patchvolume_004": "3ad54ed31e99513916f1bdf931b8694e32b524c70791d1dfba8031bc279a1ffb",
 53    "patchvolume_005": "09b860c25291b1a99f82d790cfa2b6296f50056e863ce8004e68d5c113dab666",
 54    "patchvolume_006": "839cd4b5a8ca6a54ba8c805fcfbaa03b3fa2e570d4136e497aa611275980d7e4",
 55    "patchvolume_007": "a067cc4587290745eb8964faadb8037e8a3b48ee0f471efc08ab11fbf9205ec8",
 56    "patchvolume_008": "908bf1bce5ea16da4eb8275dd39cfbec03bbb330f1e0e3fd59fdb4fbdd16ebaf",
 57    "patchvolume_009": "2bb38e298870fd783b986e3bb953a648f07dbe50b55e242b2eb2e740984d66ca",
 58    "patchvolume_010": "41f240844354753118398633278733d14ddbc04f74121f4cfe8f524761e7616f",
 59    "patchvolume_011": "3568b78e697997995123d270f3ae217659811949d3c79f9298acb2ac93c34522",
 60}
 61
 62
 63def _convert_to_h5(raw_path: str, label_path: str, h5_path: str) -> None:
 64    import h5py
 65    import nibabel as nib
 66    import numpy as np
 67
 68    if os.path.exists(h5_path):
 69        return
 70
 71    raw = np.asarray(nib.load(raw_path).dataobj)
 72    labels = np.asarray(nib.load(label_path).dataobj)
 73
 74    if raw.ndim != 3 or labels.ndim != 3 or raw.shape != labels.shape:
 75        raise RuntimeError(
 76            f"Invalid SELMA3D pair: raw shape {raw.shape} and label shape {labels.shape}."
 77        )
 78
 79    # NIfTI stores the spatial axes as XYZ. Transpose them to torch-em's ZYX convention.
 80    raw = raw.transpose(2, 1, 0).astype("float32", copy=False)
 81    labels = labels.transpose(2, 1, 0).astype("uint8", copy=False)
 82
 83    os.makedirs(os.path.dirname(h5_path), exist_ok=True)
 84    tmp_path = f"{h5_path}.incomplete"
 85    with h5py.File(tmp_path, "w") as f:
 86        f.create_dataset("raw", data=raw, chunks=(32, 128, 128), compression="gzip")
 87        f.create_dataset("label", data=labels, chunks=(32, 128, 128), compression="gzip")
 88    os.replace(tmp_path, h5_path)
 89
 90
 91def get_selma3d_data(
 92    path: Union[os.PathLike, str],
 93    split: Literal["train", "val", "test"] = "train",
 94    download: bool = False,
 95) -> str:
 96    """Download and preprocess the SELMA3D cell nucleus dataset.
 97
 98    Args:
 99        path: Filepath to a folder where the downloaded data will be saved.
100        split: The data split. One of 'train', 'val' or 'test'.
101        download: Whether to download the data if it is not present.
102
103    Returns:
104        The filepath to the preprocessed HDF5 data for the selected split.
105    """
106    if split not in SPLITS:
107        raise ValueError(f"'{split}' is not a valid split. Choose from {list(SPLITS)}.")
108
109    data_root = os.path.join(path, DATA_ROOT)
110    h5_dir = os.path.join(data_root, "h5", split)
111    expected_h5_paths = [os.path.join(h5_dir, f"{sample}.h5") for sample in SPLITS[split]]
112    if all(os.path.exists(h5_path) for h5_path in expected_h5_paths):
113        return h5_dir
114
115    raw_dir = os.path.join(data_root, "raw")
116    label_dir = os.path.join(data_root, "gt")
117    os.makedirs(raw_dir, exist_ok=True)
118    os.makedirs(label_dir, exist_ok=True)
119
120    for sample in SPLITS[split]:
121        raw_path = os.path.join(raw_dir, f"{sample}_0000.nii.gz")
122        label_path = os.path.join(label_dir, f"{sample}.nii.gz")
123        h5_path = os.path.join(h5_dir, f"{sample}.h5")
124
125        raw_url = f"{URL_ROOT}/raw/{sample}_0000.nii.gz"
126        label_url = f"{URL_ROOT}/gt/{sample}.nii.gz"
127        util.download_source(raw_path, raw_url, download, checksum=RAW_CHECKSUMS[sample])
128        util.download_source(label_path, label_url, download, checksum=LABEL_CHECKSUMS[sample])
129        _convert_to_h5(raw_path, label_path, h5_path)
130
131    return h5_dir
132
133
134def get_selma3d_paths(
135    path: Union[os.PathLike, str],
136    split: Literal["train", "val", "test"] = "train",
137    download: bool = False,
138) -> List[str]:
139    """Get paths to the SELMA3D HDF5 volumes.
140
141    Args:
142        path: Filepath to a folder where the downloaded data will be saved.
143        split: The data split. One of 'train', 'val' or 'test'.
144        download: Whether to download the data if it is not present.
145
146    Returns:
147        The filepaths to the HDF5 volumes for the selected split.
148    """
149    h5_dir = get_selma3d_data(path, split, download)
150    h5_paths = [os.path.join(h5_dir, f"{sample}.h5") for sample in SPLITS[split]]
151
152    missing_paths = [h5_path for h5_path in h5_paths if not os.path.exists(h5_path)]
153    if missing_paths:
154        raise RuntimeError(f"Could not find {len(missing_paths)} SELMA3D volumes for split '{split}'.")
155
156    return h5_paths
157
158
159def get_selma3d_dataset(
160    path: Union[os.PathLike, str],
161    patch_shape: Tuple[int, int, int],
162    split: Literal["train", "val", "test"] = "train",
163    download: bool = False,
164    **kwargs,
165) -> Dataset:
166    """Get the SELMA3D dataset for semantic nucleus segmentation.
167
168    Args:
169        path: Filepath to a folder where the downloaded data will be saved.
170        patch_shape: The 3D patch shape to use for training.
171        split: The data split. One of 'train', 'val' or 'test'.
172        download: Whether to download the data if it is not present.
173        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
174
175    Returns:
176        The segmentation dataset.
177    """
178    if len(patch_shape) != 3:
179        raise ValueError(f"The SELMA3D patch shape must be three-dimensional, got {patch_shape}.")
180
181    h5_paths = get_selma3d_paths(path, split, download)
182    return torch_em.default_segmentation_dataset(
183        raw_paths=h5_paths,
184        raw_key="raw",
185        label_paths=h5_paths,
186        label_key="label",
187        patch_shape=patch_shape,
188        ndim=3,
189        **kwargs,
190    )
191
192
193def get_selma3d_loader(
194    path: Union[os.PathLike, str],
195    batch_size: int,
196    patch_shape: Tuple[int, int, int],
197    split: Literal["train", "val", "test"] = "train",
198    download: bool = False,
199    **kwargs,
200) -> DataLoader:
201    """Get the SELMA3D dataloader for semantic nucleus segmentation.
202
203    Args:
204        path: Filepath to a folder where the downloaded data will be saved.
205        batch_size: The batch size for training.
206        patch_shape: The 3D patch shape to use for training.
207        split: The data split. One of 'train', 'val' or 'test'.
208        download: Whether to download the data if it is not present.
209        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or the PyTorch DataLoader.
210
211    Returns:
212        The DataLoader.
213    """
214    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
215    dataset = get_selma3d_dataset(
216        path=path,
217        patch_shape=patch_shape,
218        split=split,
219        download=download,
220        **ds_kwargs,
221    )
222    return torch_em.get_data_loader(dataset, batch_size=batch_size, **loader_kwargs)