torch_em.data.datasets.light_microscopy.cshaper

  1"""The CShaper dataset contains 3D fluorescence microscopy images of Caenorhabditis
  2elegans early embryos with cell instance segmentation annotations.
  3
  4The dataset is organised into training and evaluation splits:
  5- Training: Sample01, Sample02 (27 timepoints each)
  6- Evaluation: Sample02, Sample03, Sample04 (7 timepoints each)
  7
  8Each timepoint is a separate 3D NIfTI volume (.nii.gz):
  9- Raw membrane images: RawMemb/{sample}_{tp}_rawMemb.nii.gz
 10- Cell segmentation: SegCell/{sample}_{tp}_segCell.nii.gz
 11
 12NOTE: You must download the data manually. The figshare record
 13https://doi.org/10.6084/m9.figshare.12839315 links to a Google Drive folder,
 14https://drive.google.com/drive/folders/1pVhnvvliE_F7tv1zvGtjav_jWuELVbDj.
 15Download "CShaper Supplementary Data" from that folder as a zip.
 16Then place the zip inside `path`.
 17
 18The dataset is from the publication https://doi.org/10.1038/s41467-020-19863-x.
 19Please cite it if you use this dataset in your research.
 20"""
 21
 22import os
 23from glob import glob
 24from natsort import natsorted
 25from typing import List, Literal, Optional, Tuple, Union
 26
 27from torch.utils.data import Dataset, DataLoader
 28
 29import torch_em
 30
 31from .. import util
 32
 33
 34# Root path inside the zip after extraction
 35_ZIP_ROOT = "CShaper Supplementary Data/DMapNet Training and Evaluation"
 36
 37TRAIN_SAMPLES = ["Sample01", "Sample02"]
 38EVAL_SAMPLES = ["Sample02", "Sample03", "Sample04"]
 39
 40
 41def get_cshaper_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 42    """Extract the CShaper dataset zip.
 43
 44    NOTE: You must download the zip manually from the Google Drive folder that
 45    https://doi.org/10.6084/m9.figshare.12839315 links to. Place the zip inside `path`.
 46    This function extracts any zip file that it finds in `path`.
 47
 48    Args:
 49        path: Filepath to a folder containing the downloaded CShaper zip.
 50        download: Ignored (manual download required).
 51
 52    Returns:
 53        The filepath to the extracted data root directory.
 54    """
 55    data_dir = os.path.join(path, _ZIP_ROOT)
 56    if os.path.exists(data_dir):
 57        return data_dir
 58
 59    # Find any zip in path
 60    zips = glob(os.path.join(path, "*.zip"))
 61    if not zips:
 62        raise RuntimeError(
 63            f"The folder {path} contains no zip file. "
 64            "Download the CShaper data from the Google Drive folder that "
 65            "https://doi.org/10.6084/m9.figshare.12839315 links to. Then place the zip in that folder."
 66        )
 67
 68    util.unzip(zips[0], path)
 69    return data_dir
 70
 71
 72def _convert_to_h5(data_dir: str, split: str) -> str:
 73    """Convert NIfTI timepoint files to per-timepoint HDF5 files.
 74
 75    Args:
 76        data_dir: The extracted CShaper root directory.
 77        split: "train" or "val".
 78
 79    Returns:
 80        The directory containing the converted HDF5 files.
 81    """
 82    split_subdir = "TrainingData" if split == "train" else "EvaluationData"
 83    split_dir = os.path.join(data_dir, split_subdir)
 84
 85    h5_dir = os.path.join(data_dir, f"h5_{split}")
 86    if os.path.exists(h5_dir) and len(glob(os.path.join(h5_dir, "*.h5"))) > 0:
 87        return h5_dir
 88
 89    # nibabel is only needed for the conversion, so that the converted data also loads without it.
 90    try:
 91        import nibabel as nib
 92    except ImportError:
 93        raise RuntimeError(
 94            "The 'nibabel' package is required to read CShaper NIfTI files. "
 95            "Install with: pip install nibabel"
 96        )
 97    import h5py
 98
 99    os.makedirs(h5_dir, exist_ok=True)
100
101    sample_dirs = natsorted([
102        d for d in glob(os.path.join(split_dir, "*/")) if os.path.isdir(d)
103    ])
104
105    for sample_dir in sample_dirs:
106        raw_files = natsorted(glob(os.path.join(sample_dir, "RawMemb", "*.nii.gz")))
107        seg_dir = os.path.join(sample_dir, "SegCell")
108
109        for raw_path in raw_files:
110            # e.g. Sample01_030_rawMemb.nii.gz -> Sample01_030
111            basename = os.path.basename(raw_path)
112            tp_stem = basename.replace("_rawMemb.nii.gz", "")
113            h5_path = os.path.join(h5_dir, f"{tp_stem}.h5")
114
115            if os.path.exists(h5_path):
116                continue
117
118            seg_path = os.path.join(seg_dir, f"{tp_stem}_segCell.nii.gz")
119            if not os.path.exists(seg_path):
120                continue
121
122            raw_vol = nib.load(raw_path).get_fdata().astype("float32")
123            seg_vol = nib.load(seg_path).get_fdata().astype("int32")
124
125            with h5py.File(h5_path, "w") as f:
126                f.create_dataset("raw", data=raw_vol, compression="gzip")
127                f.create_dataset("labels", data=seg_vol, compression="gzip")
128
129    return h5_dir
130
131
132def get_cshaper_paths(
133    path: Union[os.PathLike, str],
134    split: Literal["train", "val"] = "train",
135    samples: Optional[List[str]] = None,
136    download: bool = False,
137) -> Tuple[List[str], List[str]]:
138    """Get paths to the CShaper data.
139
140    Args:
141        path: Filepath to a folder containing the downloaded CShaper zip.
142        split: The data split to use. Either "train" (Sample01, Sample02) or
143            "val" (Sample02, Sample03, Sample04).
144        samples: Optional list of sample names to restrict to (e.g., ["Sample01"]).
145            If None, all samples for the split are used.
146        download: Ignored (manual download required).
147
148    Returns:
149        List of filepaths for the HDF5 image data (key: "raw").
150        List of filepaths for the HDF5 label data (key: "labels").
151    """
152    if split not in ("train", "val"):
153        raise ValueError(f"Invalid split '{split}'. Choose 'train' or 'val'.")
154
155    data_dir = get_cshaper_data(path, download)
156    h5_dir = _convert_to_h5(data_dir, split)
157
158    h5_files = natsorted(glob(os.path.join(h5_dir, "*.h5")))
159
160    if len(h5_files) == 0:
161        raise RuntimeError(f"No HDF5 files found in {h5_dir}. Check the dataset structure.")
162
163    if samples is not None:
164        h5_files = [p for p in h5_files if any(os.path.basename(p).startswith(s) for s in samples)]
165
166    return h5_files, h5_files
167
168
169def get_cshaper_dataset(
170    path: Union[os.PathLike, str],
171    patch_shape: Tuple[int, ...],
172    split: Literal["train", "val"] = "train",
173    samples: Optional[List[str]] = None,
174    raw_key: str = "raw",
175    label_key: str = "labels",
176    download: bool = False,
177    **kwargs,
178) -> Dataset:
179    """Get the CShaper dataset for C. elegans embryo cell segmentation.
180
181    Args:
182        path: Filepath to a folder containing the downloaded CShaper zip.
183        patch_shape: The patch shape to use for training.
184        split: The data split to use. Either "train" or "val".
185        samples: Optional list of sample names to restrict to (e.g., ["Sample01"]).
186        raw_key: The HDF5 key for raw image data.
187        label_key: The HDF5 key for label data.
188        download: Ignored (manual download required).
189        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
190
191    Returns:
192        The segmentation dataset.
193    """
194    raw_paths, label_paths = get_cshaper_paths(path, split, samples, download)
195
196    return torch_em.default_segmentation_dataset(
197        raw_paths=raw_paths,
198        raw_key=raw_key,
199        label_paths=label_paths,
200        label_key=label_key,
201        patch_shape=patch_shape,
202        **kwargs,
203    )
204
205
206def get_cshaper_loader(
207    path: Union[os.PathLike, str],
208    batch_size: int,
209    patch_shape: Tuple[int, ...],
210    split: Literal["train", "val"] = "train",
211    samples: Optional[List[str]] = None,
212    raw_key: str = "raw",
213    label_key: str = "labels",
214    download: bool = False,
215    **kwargs,
216) -> DataLoader:
217    """Get the CShaper dataloader for C. elegans embryo cell segmentation.
218
219    Args:
220        path: Filepath to a folder containing the downloaded CShaper zip.
221        batch_size: The batch size for training.
222        patch_shape: The patch shape to use for training.
223        split: The data split to use. Either "train" or "val".
224        samples: Optional list of sample names to restrict to (e.g., ["Sample01"]).
225        raw_key: The HDF5 key for raw image data.
226        label_key: The HDF5 key for label data.
227        download: Ignored (manual download required).
228        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
229
230    Returns:
231        The DataLoader.
232    """
233    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
234    dataset = get_cshaper_dataset(path, patch_shape, split, samples, raw_key, label_key, download, **ds_kwargs)
235    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)