torch_em.data.datasets.light_microscopy.neurips_cell_seg

  1"""This dataset comes from the Neurips Cell Segmentation Challenge,
  2which collects microscopy images and annotations for cell segmentation.
  3
  4The dataset contains both images with annotations for cell segmentation
  5and unlabed images for self-supervised or semi-supervised learning.
  6See also the challenge website for details: https://neurips22-cellseg.grand-challenge.org/.
  7The dataset os decribed in the publication https://doi.org/10.1038/s41592-024-02233-6.
  8Please cite it if you use the dataset in your research.
  9"""
 10
 11import os
 12import numpy as np
 13from glob import glob
 14from typing import Union, Tuple, Any, Optional
 15
 16import torch
 17import torch_em
 18from torch.utils.data import Dataset, DataLoader
 19
 20from .. import util
 21from ... import ImageCollectionDataset, RawImageCollectionDataset, ConcatDataset
 22
 23
 24URL = {
 25    "train": "https://zenodo.org/records/10719375/files/Training-labeled.zip",
 26    "val": "https://zenodo.org/records/10719375/files/Tuning.zip",
 27    "test": "https://zenodo.org/records/10719375/files/Testing.zip",
 28    "unlabeled": "https://zenodo.org/records/10719375/files/train-unlabeled-part1.zip",
 29    "unlabeled_wsi": "https://zenodo.org/records/10719375/files/train-unlabeled-part2.zip"
 30}
 31
 32CHECKSUM = {
 33    "train": "b2383929eb8e99b2716fa0d4e2f6e03983e626a57cf00fe85175869c54aa3592",
 34    "val": "849423d36bb8fcc2d91a5b189a3b6d93c3d4071c9701eaaa44ba393a510459c4",
 35    "test": "3379730221f43830d30fddf131750e967c9c9bdf04f98811e852a050eb659ccc",
 36    "unlabeled": "390b38b398b05e9e5306a024a3bd48ab22e49592cfab3c1a119eab3636b38e0d",
 37    "unlabeled_wsi": "d1e68eba2918305eab8b846e7578ac14683de970e3fa6a7c2a4a55753be56204"
 38}
 39
 40
 41DIR_NAMES = {
 42    "train": "Training-labeled", "val": "Tuning", "test": "Testing/Public",
 43    "unlabeled": "release-part1", "unlabeled_wsi": "train-unlabeled-part2"
 44}
 45
 46ZIP_PATH = {
 47    "train": "Training-labeled.zip", "val": "Tuning.zip", "test": "Testing.zip",
 48    "unlabeled": "train-unlabeled-part1.zip", "unlabeled_wsi": "train-unlabeled-part2.zip"
 49}
 50
 51
 52def to_rgb(image):
 53    if image.ndim == 2:
 54        image = np.concatenate([image[None]] * 3, axis=0)
 55
 56    if image.ndim == 3 and image.shape[-1] == 3:
 57        image = image.transpose(2, 0, 1)
 58
 59    assert image.ndim == 3
 60    assert image.shape[0] == 3, f"{image.shape}"
 61    return image
 62
 63
 64def get_neurips_cellseg_data(root: Union[os.PathLike, str], split: str, download: bool) -> str:
 65    f"""Download the Neurips Cell Seg training data.
 66
 67    Args:
 68        root: Filepath to a folder where the downloaded data will be saved.
 69        split: The data split to download. Available splits are:
 70            {', '.join(URL.keys())}
 71        download: Whether to download the data if it is not present.
 72
 73    Returns:
 74        The filepath to the training data.
 75    """
 76    os.makedirs(root, exist_ok=True)
 77
 78    target_dir = os.path.join(root, DIR_NAMES[split])
 79    zip_path = os.path.join(root, ZIP_PATH[split])
 80
 81    if not os.path.exists(target_dir):
 82        util.download_source(path=zip_path, url=URL[split], download=download, checksum=CHECKSUM[split])
 83        util.unzip(zip_path=zip_path, dst=root)
 84
 85    return target_dir
 86
 87
 88def _get_image_and_label_paths(root, split, download):
 89    path = get_neurips_cellseg_data(root, split, download)
 90
 91    image_folder = os.path.join(path, "images")
 92    assert os.path.exists(image_folder)
 93    label_folder = os.path.join(path, "labels")
 94    assert os.path.exists(label_folder)
 95
 96    all_image_paths = glob(os.path.join(image_folder, "*"))
 97    all_image_paths.sort()
 98    all_label_paths = glob(os.path.join(label_folder, "*"))
 99    all_label_paths.sort()
100    assert len(all_image_paths) == len(all_label_paths)
101
102    return all_image_paths, all_label_paths
103
104
105def get_neurips_cellseg_supervised_dataset(
106    root: Union[str, os.PathLike],
107    split: str,
108    patch_shape: Tuple[int, int],
109    make_rgb: bool = True,
110    label_transform: Optional[Any] = None,
111    label_transform2: Optional[Any] = None,
112    raw_transform: Optional[Any] = None,
113    transform: Optional[Any] = None,
114    label_dtype: torch.dtype = torch.float32,
115    n_samples: Optional[int] = None,
116    sampler: Optional[Any] = None,
117    download: bool = False,
118) -> Dataset:
119    f"""Get the dataset for cell segmentation from the NeurIPS Cell Seg Challenge.
120
121    Args:
122        root: Filepath to a folder where the downloaded data will be saved.
123        split: The data split to download. Available splits are:
124            {', '.join(URL.keys())}
125        patch_shape: The patch shape to use for training.
126        make_rgb: Whether to map all data to RGB or treat it as grayscale.
127        label_transform: Transformation of labels, applied before data augmentation.
128        label_transform2: Transformation of labels, applied after data augmentation.
129        raw_transform: Transformation of the raw data.
130        label_dtype: The data type of the label data.
131        n_samples: Number of samples per epoch from this dataset.
132        sampler: Sampler for rejecting batches.
133        download: Whether to download the data if it is not present.
134
135    Returns:
136        The segmentation dataset.
137    """
138    assert split in ("train", "val", "test"), split
139    image_paths, label_paths = _get_image_and_label_paths(root, split, download)
140
141    if raw_transform is None:
142        trafo = to_rgb if make_rgb else None
143        raw_transform = torch_em.transform.get_raw_transform(augmentation2=trafo)
144    if transform is None:
145        transform = torch_em.transform.get_augmentations(ndim=2)
146
147    ds = ImageCollectionDataset(
148        raw_image_paths=image_paths,
149        label_image_paths=label_paths,
150        patch_shape=patch_shape,
151        raw_transform=raw_transform,
152        label_transform=label_transform,
153        label_transform2=label_transform2,
154        label_dtype=label_dtype,
155        transform=transform,
156        n_samples=n_samples,
157        sampler=sampler
158    )
159    return ds
160
161
162def get_neurips_cellseg_supervised_loader(
163    root: Union[str, os.PathLike],
164    split: str,
165    patch_shape: Tuple[int, int],
166    batch_size: int,
167    make_rgb: bool = True,
168    label_transform: Optional[Any] = None,
169    label_transform2: Optional[Any] = None,
170    raw_transform: Optional[Any] = None,
171    transform: Optional[Any] = None,
172    label_dtype: torch.dtype = torch.float32,
173    n_samples: Optional[Any] = None,
174    sampler: Optional[Any] = None,
175    download: bool = False,
176    **loader_kwargs
177) -> DataLoader:
178    f"""Get the dataset for cell segmentation from the NeurIPS Cell Seg Challenge.
179
180    Args:
181        root: Filepath to a folder where the downloaded data will be saved.
182        split: The data split to download. Available splits are:
183            {', '.join(URL.keys())}
184        patch_shape: The patch shape to use for training.
185        batch_size: The batch size for training.
186        make_rgb: Whether to map all data to RGB or treat it as grayscale.
187        label_transform: Transformation of labels, applied before data augmentation.
188        label_transform2: Transformation of labels, applied after data augmentation.
189        raw_transform: Transformation of the raw data.
190        transform: Transformation applied to raw and label data.
191        label_dtype: The data type of the label data.
192        n_samples: Number of samples per epoch from this dataset.
193        sampler: Sampler for rejecting batches.
194        download: Whether to download the data if it is not present.
195        loader_kwargs: Keyword arguments for the PyTorch DataLoader.
196
197    Returns:
198        The DataLoader.
199    """
200    ds = get_neurips_cellseg_supervised_dataset(
201        root=root,
202        split=split,
203        patch_shape=patch_shape,
204        make_rgb=make_rgb,
205        label_transform=label_transform,
206        label_transform2=label_transform2,
207        raw_transform=raw_transform,
208        transform=transform,
209        label_dtype=label_dtype,
210        n_samples=n_samples,
211        sampler=sampler,
212        download=download
213    )
214    return torch_em.segmentation.get_data_loader(ds, batch_size, **loader_kwargs)
215
216
217def _get_image_paths(root, download):
218    path = get_neurips_cellseg_data(root, "unlabeled", download)
219    image_paths = glob(os.path.join(path, "*"))
220    image_paths.sort()
221    return image_paths
222
223
224def _get_wholeslide_paths(root, patch_shape, download):
225    path = get_neurips_cellseg_data(root, "unlabeled_wsi", download)
226    image_paths = glob(os.path.join(path, "*"))
227    image_paths.sort()
228
229    # one of the whole slides doesn't support memmap which will make it very slow to load
230    image_paths = [path for path in image_paths if torch_em.util.supports_memmap(path)]
231    assert len(image_paths) > 0
232
233    n_samples = 0
234    for im_path in image_paths:
235        shape = torch_em.util.load_image(im_path).shape
236        assert len(shape) == 3 and shape[-1] == 3
237        im_shape = shape[:2]
238        n_samples += np.prod([sh // psh for sh, psh in zip(im_shape, patch_shape)])
239
240    return image_paths, n_samples
241
242
243def get_neurips_cellseg_unsupervised_dataset(
244    root: Union[str, os.PathLike],
245    patch_shape: Tuple[int, int],
246    make_rgb: bool = True,
247    raw_transform: Optional[Any] = None,
248    transform: Optional[Any] = None,
249    dtype: torch.dtype = torch.float32,
250    sampler: Optional[Any] = None,
251    use_images: bool = True,
252    use_wholeslide: bool = True,
253    download: bool = False,
254) -> Dataset:
255    """Get the unsupervised dataset from the NeurIPS Cell Seg Challenge.
256
257    Args:
258        root: Filepath to a folder where the downloaded data will be saved.
259        patch_shape: The patch shape to use for training.
260        make_rgb: Whether to map all data to RGB or treat it as grayscale.
261        raw_transform: Transformation of the raw data.
262        transform: Transformation applied to raw and label data.
263        dtype: The data type of the image data.
264        sampler: Sampler for rejecting batches.
265        use_images: Whether to use the normal image data.
266        use_wholeslide: Whether to use the wholeslide image data.
267        download: Whether to download the data if it is not present.
268
269    Returns:
270        The segmentation dataset.
271    """
272    if raw_transform is None:
273        trafo = to_rgb if make_rgb else None
274        raw_transform = torch_em.transform.get_raw_transform(augmentation2=trafo)
275    if transform is None:
276        transform = torch_em.transform.get_augmentations(ndim=2)
277
278    datasets = []
279    if use_images:
280        image_paths = _get_image_paths(root, download)
281        datasets.append(
282            RawImageCollectionDataset(
283                raw_image_paths=image_paths,
284                patch_shape=patch_shape,
285                raw_transform=raw_transform,
286                transform=transform,
287                dtype=dtype,
288                sampler=sampler
289            )
290        )
291    if use_wholeslide:
292        image_paths, n_samples = _get_wholeslide_paths(root, patch_shape, download)
293        datasets.append(
294            RawImageCollectionDataset(
295                raw_image_paths=image_paths,
296                patch_shape=patch_shape,
297                raw_transform=raw_transform,
298                transform=transform,
299                dtype=dtype,
300                n_samples=n_samples,
301                sampler=sampler
302            )
303        )
304    assert len(datasets) > 0
305    return ConcatDataset(*datasets)
306
307
308def get_neurips_cellseg_unsupervised_loader(
309    root: Union[str, os.PathLike],
310    patch_shape: Tuple[int, int],
311    batch_size: int,
312    make_rgb: bool = True,
313    raw_transform: Optional[Any] = None,
314    transform: Optional[Any] = None,
315    dtype: torch.dtype = torch.float32,
316    sampler: Optional[Any] = None,
317    use_images: bool = True,
318    use_wholeslide: bool = True,
319    download: bool = False,
320    **loader_kwargs,
321) -> DataLoader:
322    """Get the unsupervised dataset from the NeurIPS Cell Seg Challenge.
323
324    Args:
325        root: Filepath to a folder where the downloaded data will be saved.
326        patch_shape: The patch shape to use for training.
327        batch_size: The batch size for training.
328        make_rgb: Whether to map all data to RGB or treat it as grayscale.
329        raw_transform: Transformation of the raw data.
330        transform: Transformation applied to raw and label data.
331        dtype: The data type of the image data.
332        sampler: Sampler for rejecting batches.
333        use_images: Whether to use the normal image data.
334        use_wholeslide: Whether to use the wholeslide image data.
335        download: Whether to download the data if it is not present.
336        loader_kwargs: Keyword arguments for the PyTorch DataLoader.
337
338    Returns:
339        The DataLoader.
340    """
341    ds = get_neurips_cellseg_unsupervised_dataset(
342        root=root, patch_shape=patch_shape, make_rgb=make_rgb, raw_transform=raw_transform, transform=transform,
343        dtype=dtype, sampler=sampler, use_images=use_images, use_wholeslide=use_wholeslide, download=download
344    )
345    return torch_em.segmentation.get_data_loader(ds, batch_size, **loader_kwargs)