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    try:
 83        import nibabel as nib
 84    except ImportError:
 85        raise RuntimeError(
 86            "The 'nibabel' package is required to read CShaper NIfTI files. "
 87            "Install with: pip install nibabel"
 88        )
 89    import h5py
 90
 91    split_subdir = "TrainingData" if split == "train" else "EvaluationData"
 92    split_dir = os.path.join(data_dir, split_subdir)
 93
 94    h5_dir = os.path.join(data_dir, f"h5_{split}")
 95    if os.path.exists(h5_dir) and len(glob(os.path.join(h5_dir, "*.h5"))) > 0:
 96        return h5_dir
 97    os.makedirs(h5_dir, exist_ok=True)
 98
 99    sample_dirs = natsorted([
100        d for d in glob(os.path.join(split_dir, "*/")) if os.path.isdir(d)
101    ])
102
103    for sample_dir in sample_dirs:
104        raw_files = natsorted(glob(os.path.join(sample_dir, "RawMemb", "*.nii.gz")))
105        seg_dir = os.path.join(sample_dir, "SegCell")
106
107        for raw_path in raw_files:
108            # e.g. Sample01_030_rawMemb.nii.gz -> Sample01_030
109            basename = os.path.basename(raw_path)
110            tp_stem = basename.replace("_rawMemb.nii.gz", "")
111            h5_path = os.path.join(h5_dir, f"{tp_stem}.h5")
112
113            if os.path.exists(h5_path):
114                continue
115
116            seg_path = os.path.join(seg_dir, f"{tp_stem}_segCell.nii.gz")
117            if not os.path.exists(seg_path):
118                continue
119
120            raw_vol = nib.load(raw_path).get_fdata().astype("float32")
121            seg_vol = nib.load(seg_path).get_fdata().astype("int32")
122
123            with h5py.File(h5_path, "w") as f:
124                f.create_dataset("raw", data=raw_vol, compression="gzip")
125                f.create_dataset("labels", data=seg_vol, compression="gzip")
126
127    return h5_dir
128
129
130def get_cshaper_paths(
131    path: Union[os.PathLike, str],
132    split: Literal["train", "val"] = "train",
133    samples: Optional[List[str]] = None,
134    download: bool = False,
135) -> Tuple[List[str], List[str]]:
136    """Get paths to the CShaper data.
137
138    Args:
139        path: Filepath to a folder containing the downloaded CShaper zip.
140        split: The data split to use. Either "train" (Sample01, Sample02) or
141            "val" (Sample02, Sample03, Sample04).
142        samples: Optional list of sample names to restrict to (e.g., ["Sample01"]).
143            If None, all samples for the split are used.
144        download: Ignored (manual download required).
145
146    Returns:
147        List of filepaths for the HDF5 image data (key: "raw").
148        List of filepaths for the HDF5 label data (key: "labels").
149    """
150    if split not in ("train", "val"):
151        raise ValueError(f"Invalid split '{split}'. Choose 'train' or 'val'.")
152
153    data_dir = get_cshaper_data(path, download)
154    h5_dir = _convert_to_h5(data_dir, split)
155
156    h5_files = natsorted(glob(os.path.join(h5_dir, "*.h5")))
157
158    if len(h5_files) == 0:
159        raise RuntimeError(f"No HDF5 files found in {h5_dir}. Check the dataset structure.")
160
161    if samples is not None:
162        h5_files = [p for p in h5_files if any(os.path.basename(p).startswith(s) for s in samples)]
163
164    return h5_files, h5_files
165
166
167def get_cshaper_dataset(
168    path: Union[os.PathLike, str],
169    patch_shape: Tuple[int, ...],
170    split: Literal["train", "val"] = "train",
171    samples: Optional[List[str]] = None,
172    raw_key: str = "raw",
173    label_key: str = "labels",
174    download: bool = False,
175    **kwargs,
176) -> Dataset:
177    """Get the CShaper dataset for C. elegans embryo cell segmentation.
178
179    Args:
180        path: Filepath to a folder containing the downloaded CShaper zip.
181        patch_shape: The patch shape to use for training.
182        split: The data split to use. Either "train" or "val".
183        samples: Optional list of sample names to restrict to (e.g., ["Sample01"]).
184        raw_key: The HDF5 key for raw image data.
185        label_key: The HDF5 key for label data.
186        download: Ignored (manual download required).
187        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
188
189    Returns:
190        The segmentation dataset.
191    """
192    raw_paths, label_paths = get_cshaper_paths(path, split, samples, download)
193
194    return torch_em.default_segmentation_dataset(
195        raw_paths=raw_paths,
196        raw_key=raw_key,
197        label_paths=label_paths,
198        label_key=label_key,
199        patch_shape=patch_shape,
200        **kwargs,
201    )
202
203
204def get_cshaper_loader(
205    path: Union[os.PathLike, str],
206    batch_size: int,
207    patch_shape: Tuple[int, ...],
208    split: Literal["train", "val"] = "train",
209    samples: Optional[List[str]] = None,
210    raw_key: str = "raw",
211    label_key: str = "labels",
212    download: bool = False,
213    **kwargs,
214) -> DataLoader:
215    """Get the CShaper dataloader for C. elegans embryo cell segmentation.
216
217    Args:
218        path: Filepath to a folder containing the downloaded CShaper zip.
219        batch_size: The batch size for training.
220        patch_shape: The patch shape to use for training.
221        split: The data split to use. Either "train" or "val".
222        samples: Optional list of sample names to restrict to (e.g., ["Sample01"]).
223        raw_key: The HDF5 key for raw image data.
224        label_key: The HDF5 key for label data.
225        download: Ignored (manual download required).
226        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
227
228    Returns:
229        The DataLoader.
230    """
231    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
232    dataset = get_cshaper_dataset(path, patch_shape, split, samples, raw_key, label_key, download, **ds_kwargs)
233    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)