torch_em.data.datasets.light_microscopy.neurosphere

  1"""The Neurosphere dataset contains a 3D fluorescence light-sheet microscopy image
  2of a cancer cell neurosphere with ground truth instance segmentation from the
  3OpenSegSPIM analysis pipeline.
  4
  5The dataset consists of a single volume of approximately 115 x 150 x 150 voxels.
  6
  7NOTE: The segmentations are pixelated at the boundaries and don't exactly match the segmentation.
  8
  9The dataset is located at https://sourceforge.net/projects/opensegspim/.
 10This dataset is from the publication https://doi.org/10.1093/bioinformatics/btw093.
 11Please cite it if you use this dataset in your research.
 12"""
 13
 14import os
 15from glob import glob
 16from natsort import natsorted
 17from typing import List, Tuple, Union
 18
 19import numpy as np
 20from scipy.ndimage import binary_fill_holes
 21
 22from torch.utils.data import Dataset, DataLoader
 23
 24import torch_em
 25
 26from .. import util
 27
 28
 29RAW_URL = "https://downloads.sourceforge.net/project/opensegspim/Sample%20Data/Neurosphere_Dataset.zip"
 30LABEL_URL = "https://downloads.sourceforge.net/project/opensegspim/Sample%20Data/Neurosphere_OpenSegSPIM.zip"
 31RAW_CHECKSUM = None
 32LABEL_CHECKSUM = None
 33
 34
 35def get_neurosphere_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 36    """Download the Neurosphere dataset.
 37
 38    Args:
 39        path: Filepath to a folder where the downloaded data will be saved.
 40        download: Whether to download the data if it is not present.
 41
 42    Returns:
 43        The filepath to the extracted data directory.
 44    """
 45    data_dir = os.path.join(path, "neurosphere")
 46    if os.path.exists(data_dir):
 47        return data_dir
 48
 49    os.makedirs(data_dir, exist_ok=True)
 50
 51    raw_zip = os.path.join(path, "Neurosphere_Dataset.zip")
 52    label_zip = os.path.join(path, "Neurosphere_OpenSegSPIM.zip")
 53
 54    util.download_source(raw_zip, RAW_URL, download, checksum=RAW_CHECKSUM)
 55    util.download_source(label_zip, LABEL_URL, download, checksum=LABEL_CHECKSUM)
 56
 57    util.unzip(raw_zip, data_dir)
 58    util.unzip(label_zip, data_dir)
 59
 60    return data_dir
 61
 62
 63def _fill_labels(data_dir: str) -> str:
 64    """Convert thin-shell contour labels to filled 3D instance segmentations.
 65
 66    Loads Nucleisegmented2.tif, applies binary_fill_holes per instance,
 67    renumbers to sequential IDs (1, 2, 3 ...), and saves as filled_labels.tif.
 68
 69    Args:
 70        data_dir: The neurosphere data directory.
 71
 72    Returns:
 73        Path to the filled label file.
 74    """
 75    import imageio.v3 as imageio
 76
 77    filled_path = os.path.join(data_dir, "filled_labels.tif")
 78    if os.path.exists(filled_path):
 79        return filled_path
 80
 81    label_paths = natsorted(glob(os.path.join(data_dir, "**", "Nucleisegmented2.tif"), recursive=True))
 82    if len(label_paths) == 0:
 83        raise RuntimeError(f"Label file 'Nucleisegmented2.tif' not found in {data_dir}.")
 84
 85    raw_labels = imageio.imread(label_paths[0])
 86    instance_ids = np.unique(raw_labels)
 87    instance_ids = instance_ids[instance_ids != 0]
 88
 89    filled = np.zeros(raw_labels.shape, dtype=np.int32)
 90    for new_id, val in enumerate(instance_ids, start=1):
 91        mask = binary_fill_holes(raw_labels == val)
 92        filled[mask] = new_id
 93
 94    imageio.imwrite(filled_path, filled)
 95    return filled_path
 96
 97
 98def get_neurosphere_paths(
 99    path: Union[os.PathLike, str], download: bool = False,
100) -> Tuple[List[str], List[str]]:
101    """Get paths to the Neurosphere data.
102
103    Args:
104        path: Filepath to a folder where the downloaded data will be saved.
105        download: Whether to download the data if it is not present.
106
107    Returns:
108        List of filepaths for the image data.
109        List of filepaths for the label data.
110    """
111    data_dir = get_neurosphere_data(path, download)
112
113    raw_paths = natsorted(glob(os.path.join(data_dir, "OriginalStack.tif")))
114    if len(raw_paths) == 0:
115        raise RuntimeError(
116            f"Raw image 'OriginalStack.tif' not found in {data_dir}. "
117            "Please check the dataset structure after downloading."
118        )
119
120    filled_label_path = _fill_labels(data_dir)
121    label_paths = [filled_label_path]
122
123    return raw_paths, label_paths
124
125
126def get_neurosphere_dataset(
127    path: Union[os.PathLike, str],
128    patch_shape: Tuple[int, ...],
129    download: bool = False,
130    **kwargs,
131) -> Dataset:
132    """Get the Neurosphere dataset for 3D cell instance segmentation.
133
134    Args:
135        path: Filepath to a folder where the downloaded data will be saved.
136        patch_shape: The patch shape to use for training.
137        download: Whether to download the data if it is not present.
138        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
139
140    Returns:
141        The segmentation dataset.
142    """
143    raw_paths, label_paths = get_neurosphere_paths(path, download)
144
145    return torch_em.default_segmentation_dataset(
146        raw_paths=raw_paths,
147        raw_key=None,
148        label_paths=label_paths,
149        label_key=None,
150        patch_shape=patch_shape,
151        **kwargs,
152    )
153
154
155def get_neurosphere_loader(
156    path: Union[os.PathLike, str],
157    batch_size: int,
158    patch_shape: Tuple[int, ...],
159    download: bool = False,
160    **kwargs,
161) -> DataLoader:
162    """Get the Neurosphere dataloader for 3D cell instance segmentation.
163
164    Args:
165        path: Filepath to a folder where the downloaded data will be saved.
166        batch_size: The batch size for training.
167        patch_shape: The patch shape to use for training.
168        download: Whether to download the data if it is not present.
169        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
170
171    Returns:
172        The DataLoader.
173    """
174    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
175    dataset = get_neurosphere_dataset(path, patch_shape, download, **ds_kwargs)
176    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)