torch_em.data.datasets.light_microscopy.vampire

  1"""The VAMPIRE dataset contains annotations for cell segmentation in
  2fluorescence microscopy images of mouse embryonic fibroblasts.
  3
  4The cells are stained with phalloidin and imaged in a single channel at 1024 x 1280 pixels.
  5They come in two conditions, wildtype cells and cells with a Lamin A/C knockout, with 30 images
  6each. The masks were produced with the CellProfiler pipeline that the authors ship next to the
  7images, and they label whole cells.
  8
  9NOTE: The annotations are sparse. The CellProfiler pipeline drops objects that touch the image
 10border and objects outside its size filters, so roughly a third of the cells in the wildtype images
 11and roughly half of the cells in the knockout images carry no mask. Unlabeled cells are part of the
 12background, so a sampler or a loss that ignores the background is advisable.
 13
 14NOTE: The study deposits two further image sets that carry no masks and are therefore not part of
 15this loader. Mouse embryonic fibroblasts on micropatterns are at
 16https://github.com/kukionfr/Micropattern_MEF_LMNA_Image and human dermal fibroblast nuclei of seven
 17donor ages are at https://github.com/kukionfr/Aging_human_dermal_fibroblast_nucleus . Both hold
 18morphometric tables of the cells rather than segmentations.
 19
 20The data is published as the supplementary data of the VAMPIRE software at
 21https://github.com/kukionfr/VAMPIRE_open . This dataset is from the publication
 22https://doi.org/10.1038/s41596-020-00432-x .
 23Please cite it if you use this dataset in your research.
 24"""
 25
 26import os
 27from glob import glob
 28from natsort import natsorted
 29from typing import List, Tuple, Union, Literal, Optional
 30
 31from torch.utils.data import Dataset, DataLoader
 32
 33import torch_em
 34
 35from .. import util
 36
 37
 38URL = "https://github.com/kukionfr/VAMPIRE_open/releases/download/v1.0/Supplementary.Data.zip"
 39CHECKSUM = "a5e9b70537d5add8b860fc8ac4b40c7f7e260dfd21437926562c41a6d8f95da7"
 40
 41# The folder names of the two conditions in the archive.
 42SAMPLE_TYPES = {"wildtype": "MEF_wildtype", "lmna_knockout": "MEF_LMNA--"}
 43
 44
 45def _preprocess_labels(data_dir: str) -> str:
 46    """Map the masks to consecutive instance ids.
 47
 48    The archive stores the instance ids spread over the full uint16 range, so that the masks display
 49    well in an image viewer. This restores the ids that the loader expects.
 50    """
 51    import numpy as np
 52    import imageio.v3 as imageio
 53    from tqdm import tqdm
 54
 55    output_dir = os.path.join(data_dir, "preprocessed")
 56
 57    for sample_type, folder in SAMPLE_TYPES.items():
 58        label_dir = os.path.join(output_dir, sample_type)
 59        os.makedirs(label_dir, exist_ok=True)
 60
 61        label_paths = natsorted(glob(os.path.join(data_dir, "Example segmented images", folder, "*.tiff")))
 62        if not label_paths:
 63            raise RuntimeError(f"Could not find any masks for '{sample_type}' in {data_dir}.")
 64
 65        for label_path in tqdm(label_paths, desc=f"Preprocess '{sample_type}'"):
 66            output_path = os.path.join(label_dir, f"{os.path.splitext(os.path.basename(label_path))[0]}.tif")
 67            if os.path.exists(output_path):
 68                continue
 69
 70            labels = imageio.imread(label_path)
 71            ids = np.unique(labels)
 72            relabeled = np.searchsorted(ids, labels).astype("uint16")
 73            if ids[0] != 0:  # Guard against a mask without background.
 74                relabeled += 1
 75
 76            temporary_path = f"{output_path}.tmp.tif"
 77            imageio.imwrite(temporary_path, relabeled, compression="zlib")
 78            os.replace(temporary_path, output_path)
 79
 80    return output_dir
 81
 82
 83def get_vampire_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 84    """Download the VAMPIRE dataset.
 85
 86    Args:
 87        path: Filepath to a folder where the downloaded data will be saved.
 88        download: Whether to download the data if it is not present.
 89
 90    Returns:
 91        The filepath to the extracted data.
 92    """
 93    data_dir = os.path.join(path, "Supplementary Data")
 94    if os.path.exists(data_dir):
 95        return data_dir
 96
 97    os.makedirs(path, exist_ok=True)
 98    zip_path = os.path.join(path, "Supplementary.Data.zip")
 99    util.download_source(path=zip_path, url=URL, download=download, checksum=CHECKSUM)
100    util.unzip(zip_path=zip_path, dst=data_dir)
101
102    return data_dir
103
104
105def get_vampire_paths(
106    path: Union[os.PathLike, str],
107    sample_type: Optional[Literal["wildtype", "lmna_knockout"]] = None,
108    download: bool = False,
109) -> Tuple[List[str], List[str]]:
110    """Get paths to the VAMPIRE data.
111
112    Args:
113        path: Filepath to a folder where the downloaded data will be saved.
114        sample_type: The condition of the cells. Either 'wildtype' or 'lmna_knockout'. By default, both are used.
115        download: Whether to download the data if it is not present.
116
117    Returns:
118        List of filepaths to the image data.
119        List of filepaths to the label data.
120    """
121    if sample_type is None:
122        sample_types = list(SAMPLE_TYPES)
123    elif sample_type in SAMPLE_TYPES:
124        sample_types = [sample_type]
125    else:
126        raise ValueError(f"'{sample_type}' is not a valid sample type. Choose from {list(SAMPLE_TYPES)}.")
127
128    data_dir = get_vampire_data(path, download)
129    output_dir = _preprocess_labels(data_dir)
130
131    image_paths, label_paths = [], []
132    for name in sample_types:
133        for image_path in natsorted(glob(os.path.join(data_dir, "Example images", SAMPLE_TYPES[name], "*.tif"))):
134            # The archive gives an image and its mask the same file name.
135            label_path = os.path.join(output_dir, name, os.path.basename(image_path))
136            if not os.path.exists(label_path):
137                raise RuntimeError(f"Could not find the mask for the image '{image_path}' at {label_path}.")
138            image_paths.append(image_path)
139            label_paths.append(label_path)
140
141    if not image_paths:
142        raise RuntimeError(f"Could not find any VAMPIRE data in {data_dir}.")
143
144    return image_paths, label_paths
145
146
147def get_vampire_dataset(
148    path: Union[os.PathLike, str],
149    patch_shape: Tuple[int, int],
150    sample_type: Optional[Literal["wildtype", "lmna_knockout"]] = None,
151    offsets: Optional[List[List[int]]] = None,
152    boundaries: bool = False,
153    binary: bool = False,
154    download: bool = False,
155    **kwargs,
156) -> Dataset:
157    """Get the VAMPIRE dataset for cell segmentation.
158
159    Args:
160        path: Filepath to a folder where the downloaded data will be saved.
161        patch_shape: The patch shape to use for training.
162        sample_type: The condition of the cells. Either 'wildtype' or 'lmna_knockout'. By default, both are used.
163        offsets: Offset values for affinity computation used as target.
164        boundaries: Whether to compute boundaries as the target.
165        binary: Whether to use a binary segmentation target.
166        download: Whether to download the data if it is not present.
167        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
168
169    Returns:
170        The segmentation dataset.
171    """
172    image_paths, label_paths = get_vampire_paths(path, sample_type, download)
173
174    kwargs, _ = util.add_instance_label_transform(
175        kwargs, add_binary_target=True, offsets=offsets, boundaries=boundaries, binary=binary,
176    )
177
178    return torch_em.default_segmentation_dataset(
179        raw_paths=image_paths,
180        raw_key=None,
181        label_paths=label_paths,
182        label_key=None,
183        patch_shape=patch_shape,
184        is_seg_dataset=False,
185        **kwargs,
186    )
187
188
189def get_vampire_loader(
190    path: Union[os.PathLike, str],
191    batch_size: int,
192    patch_shape: Tuple[int, int],
193    sample_type: Optional[Literal["wildtype", "lmna_knockout"]] = None,
194    offsets: Optional[List[List[int]]] = None,
195    boundaries: bool = False,
196    binary: bool = False,
197    download: bool = False,
198    **kwargs,
199) -> DataLoader:
200    """Get the VAMPIRE dataloader for cell segmentation.
201
202    Args:
203        path: Filepath to a folder where the downloaded data will be saved.
204        batch_size: The batch size for training.
205        patch_shape: The patch shape to use for training.
206        sample_type: The condition of the cells. Either 'wildtype' or 'lmna_knockout'. By default, both are used.
207        offsets: Offset values for affinity computation used as target.
208        boundaries: Whether to compute boundaries as the target.
209        binary: Whether to use a binary segmentation target.
210        download: Whether to download the data if it is not present.
211        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or the PyTorch DataLoader.
212
213    Returns:
214        The DataLoader.
215    """
216    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
217    dataset = get_vampire_dataset(
218        path=path,
219        patch_shape=patch_shape,
220        sample_type=sample_type,
221        offsets=offsets,
222        boundaries=boundaries,
223        binary=binary,
224        download=download,
225        **ds_kwargs,
226    )
227    return torch_em.get_data_loader(dataset, batch_size=batch_size, **loader_kwargs)