torch_em.data.datasets.light_microscopy.covid_if

  1"""This dataset contains annotation for cell and nucleus segmentation
  2in immunofluorescence microscopy.
  3
  4This dataset is from the publication https://doi.org/10.1002/bies.202000257.
  5Please cite it if you use this dataset in your research.
  6"""
  7
  8import os
  9from glob import glob
 10from typing import List, Optional, Tuple, Union
 11
 12import torch_em
 13from torch.utils.data import Dataset, DataLoader
 14from .. import util
 15
 16COVID_IF_URL = "https://zenodo.org/record/5092850/files/covid-if-groundtruth.zip?download=1"
 17CHECKSUM = "d9cd6c85a19b802c771fb4ff928894b19a8fab0e0af269c49235fdac3f7a60e1"
 18
 19
 20def get_covid_if_data(path: Union[os.PathLike, str], download: bool) -> str:
 21    """Download the CovidIF training data.
 22
 23    Args:
 24        path: Filepath to a folder where the downloaded data will be saved.
 25        download: Whether to download the data if it is not present.
 26
 27    Returns:
 28        The filepath to the training data.
 29    """
 30    url = COVID_IF_URL
 31    checksum = CHECKSUM
 32
 33    if os.path.exists(path):
 34        return path
 35
 36    os.makedirs(path, exist_ok=True)
 37    zip_path = os.path.join(path, "covid-if.zip")
 38    util.download_source(zip_path, url, download, checksum)
 39    util.unzip(zip_path, path, True)
 40
 41    return path
 42
 43
 44def get_covid_if_dataset(
 45    path: Union[os.PathLike, str],
 46    patch_shape: Tuple[int, int],
 47    sample_range: Optional[Tuple[int, int]] = None,
 48    target: str = "cells",
 49    download: bool = False,
 50    offsets: Optional[List[List[int]]] = None,
 51    boundaries: bool = False,
 52    binary: bool = False,
 53    **kwargs
 54) -> Dataset:
 55    """Get the CovidIF dataset for segmenting nuclei or cells in immunofluorescence microscopy.
 56
 57    Args:
 58        path: Filepath to a folder where the downloaded data will be saved.
 59        patch_shape: The patch shape to use for training.
 60        sample_range: Id range of samples to load from the training dataset.
 61        target: The segmentation task. Either 'cells' or 'nuclei'.
 62        download: Whether to download the data if it is not present.
 63        offsets: Offset values for affinity computation used as target.
 64        boundaries: Whether to compute boundaries as the target.
 65        binary: Whether to use a binary segmentation target.
 66        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
 67
 68    Returns:
 69       The segmentation dataset.
 70    """
 71    available_targets = ("cells", "nuclei")
 72    # TODO also support infected_cells
 73    # available_targets = ("cells", "nuclei", "infected_cells")
 74    assert target in available_targets, f"{target} not found in {available_targets}"
 75
 76    if target == "cells":
 77        raw_key = "raw/serum_IgG/s0"
 78        label_key = "labels/cells/s0"
 79    elif target == "nuclei":
 80        raw_key = "raw/nuclei/s0"
 81        label_key = "labels/nuclei/s0"
 82
 83    get_covid_if_data(path, download)
 84
 85    file_paths = sorted(glob(os.path.join(path, "*.h5")))
 86    if sample_range is not None:
 87        start, stop = sample_range
 88        if start is None:
 89            start = 0
 90        if stop is None:
 91            stop = len(file_paths)
 92        file_paths = [os.path.join(path, f"gt_image_{idx:03}.h5") for idx in range(start, stop)]
 93        assert all(os.path.exists(fp) for fp in file_paths), f"Invalid sample range {sample_range}"
 94
 95    kwargs, _ = util.add_instance_label_transform(
 96        kwargs, add_binary_target=True, binary=binary, boundaries=boundaries, offsets=offsets
 97    )
 98    kwargs = util.update_kwargs(kwargs, "ndim", 2)
 99
100    return torch_em.default_segmentation_dataset(
101        file_paths, raw_key, file_paths, label_key, patch_shape, **kwargs
102    )
103
104
105def get_covid_if_loader(
106    path: Union[os.PathLike, str],
107    patch_shape: Tuple[int, int],
108    batch_size: int,
109    sample_range: Optional[Tuple[int, int]] = None,
110    target: str = "cells",
111    download: bool = False,
112    offsets: Optional[List[List[int]]] = None,
113    boundaries: bool = False,
114    binary: bool = False,
115    **kwargs
116) -> DataLoader:
117    """Get the CovidIF dataloder for segmenting nuclei or cells in immunofluorescence microscopy.
118
119    Args:
120        path: Filepath to a folder where the downloaded data will be saved.
121        patch_shape: The patch shape to use for training.
122        batch_size: The batch size for training.
123        sample_range: Id range of samples to load from the training dataset.
124        target: The segmentation task. Either 'cells' or 'nuclei'.
125        download: Whether to download the data if it is not present.
126        offsets: Offset values for affinity computation used as target.
127        boundaries: Whether to compute boundaries as the target.
128        binary: Whether to use a binary segmentation target.
129        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
130
131    Returns:
132        The DataLoader.
133    """
134    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
135    dataset = get_covid_if_dataset(
136        path, patch_shape, sample_range=sample_range, target=target, download=download,
137        offsets=offsets, boundaries=boundaries, binary=binary, **ds_kwargs,
138    )
139    loader = torch_em.get_data_loader(dataset, batch_size=batch_size, **loader_kwargs)
140    return loader