torch_em.data.datasets.electron_microscopy.microsam_em_nuclei

  1"""This dataset contains nucleus instance segmentations in volume electron microscopy (EM) data.
  2
  3It was generated for the user study on data annotation with microSAM
  4(https://github.com/computational-cell-analytics/micro-sam).
  5It has two 3D EM crops of size 32 x 512 x 512 voxels:
  6- 'train': The crop that was used to train the finetuned microSAM model, with manually annotated nucleus labels.
  7- 'annotation': The crop that was used for the user study, with nucleus annotations obtained via
  8  ilastik carving ('ilastik_carving'), microSAM with the default model ('default_sam')
  9  or microSAM with a finetuned model ('finetuned_sam').
 10
 11The data is available at https://doi.org/10.5281/zenodo.14037020 (CC-BY-4.0).
 12The record does not list a publication. Please cite the Zenodo record if you use this dataset.
 13"""
 14
 15import os
 16from typing import List, Literal, Optional, Tuple, Union
 17
 18from torch.utils.data import Dataset, DataLoader
 19
 20import torch_em
 21
 22from .. import util
 23
 24
 25URL = "https://zenodo.org/api/records/14037021/files/3d-user-study-data.zip/content"
 26CHECKSUM = "5eaa4650063a02becb46b6b11c1557362d3a9a9cb01fb88b3e6dc245819d4a5f"
 27
 28ANNOTATIONS = ("ilastik_carving", "default_sam", "finetuned_sam")
 29
 30
 31def get_microsam_em_nuclei_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 32    """Download the microSAM EM nuclei data.
 33
 34    Args:
 35        path: Filepath to a folder where the downloaded data will be saved.
 36        download: Whether to download the data if it is not present.
 37
 38    Returns:
 39        The filepath to the folder with the data.
 40    """
 41    data_dir = os.path.join(path, "data")
 42    if os.path.exists(data_dir):
 43        return data_dir
 44
 45    os.makedirs(path, exist_ok=True)
 46    zip_path = os.path.join(path, "3d-user-study-data.zip")
 47    util.download_source(zip_path, URL, download, checksum=CHECKSUM)
 48    util.unzip(zip_path, data_dir, remove=True)
 49    return data_dir
 50
 51
 52def get_microsam_em_nuclei_paths(
 53    path: Union[os.PathLike, str], source: str = "train", download: bool = False
 54) -> Tuple[str, Optional[str], str, Optional[str]]:
 55    """Get paths to the microSAM EM nuclei data.
 56
 57    Args:
 58        path: Filepath to a folder where the downloaded data will be saved.
 59        source: The data to use. Either 'train' (the training crop) or the annotation method
 60            of the user study crop, one of 'ilastik_carving', 'default_sam' or 'finetuned_sam'.
 61        download: Whether to download the data if it is not present.
 62
 63    Returns:
 64        The filepath to the raw data.
 65        The key for the raw data (None for tif files).
 66        The filepath to the labels.
 67        The key for the labels (None for tif files).
 68    """
 69    data_dir = get_microsam_em_nuclei_data(path, download)
 70    if source == "train":
 71        h5_path = os.path.join(data_dir, "train_data.h5")
 72        return h5_path, "raw", h5_path, "labels/nuclei"
 73    if source in ANNOTATIONS:
 74        raw_path = os.path.join(data_dir, "annotation_volume.tif")
 75        label_path = os.path.join(data_dir, "annotations", f"annotations_{source}.tif")
 76        return raw_path, None, label_path, None
 77    raise ValueError(f"'{source}' is not a valid source. Choose from 'train' or {ANNOTATIONS}.")
 78
 79
 80def get_microsam_em_nuclei_dataset(
 81    path: Union[os.PathLike, str],
 82    patch_shape: Tuple[int, int, int],
 83    source: Literal["train", "ilastik_carving", "default_sam", "finetuned_sam"] = "train",
 84    download: bool = False,
 85    offsets: Optional[List[List[int]]] = None,
 86    boundaries: bool = False,
 87    binary: bool = False,
 88    **kwargs
 89) -> Dataset:
 90    """Get the microSAM EM nuclei dataset for nucleus segmentation in EM.
 91
 92    Args:
 93        path: Filepath to a folder where the downloaded data will be saved.
 94        patch_shape: The patch shape to use for training.
 95        source: The data to use. Either 'train' (the training crop) or the annotation method
 96            of the user study crop, one of 'ilastik_carving', 'default_sam' or 'finetuned_sam'.
 97        download: Whether to download the data if it is not present.
 98        offsets: Offset values for affinity computation used as target.
 99        boundaries: Whether to compute boundaries as the target.
100        binary: Whether to return a binary segmentation target.
101        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
102
103    Returns:
104        The segmentation dataset.
105    """
106    raw_path, raw_key, label_path, label_key = get_microsam_em_nuclei_paths(path, source, download)
107
108    kwargs = util.update_kwargs(kwargs, "is_seg_dataset", True)
109    kwargs, _ = util.add_instance_label_transform(
110        kwargs, add_binary_target=True, binary=binary, boundaries=boundaries, offsets=offsets
111    )
112
113    return torch_em.default_segmentation_dataset(
114        raw_paths=raw_path, raw_key=raw_key, label_paths=label_path, label_key=label_key,
115        patch_shape=patch_shape, **kwargs
116    )
117
118
119def get_microsam_em_nuclei_loader(
120    path: Union[os.PathLike, str],
121    batch_size: int,
122    patch_shape: Tuple[int, int, int],
123    source: Literal["train", "ilastik_carving", "default_sam", "finetuned_sam"] = "train",
124    download: bool = False,
125    offsets: Optional[List[List[int]]] = None,
126    boundaries: bool = False,
127    binary: bool = False,
128    **kwargs
129) -> DataLoader:
130    """Get the microSAM EM nuclei dataloader for nucleus segmentation in EM.
131
132    Args:
133        path: Filepath to a folder where the downloaded data will be saved.
134        batch_size: The batch size for training.
135        patch_shape: The patch shape to use for training.
136        source: The data to use. Either 'train' (the training crop) or the annotation method
137            of the user study crop, one of 'ilastik_carving', 'default_sam' or 'finetuned_sam'.
138        download: Whether to download the data if it is not present.
139        offsets: Offset values for affinity computation used as target.
140        boundaries: Whether to compute boundaries as the target.
141        binary: Whether to return a binary segmentation target.
142        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
143
144    Returns:
145        The DataLoader.
146    """
147    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
148    dataset = get_microsam_em_nuclei_dataset(
149        path=path, patch_shape=patch_shape, source=source, download=download,
150        offsets=offsets, boundaries=boundaries, binary=binary, **ds_kwargs,
151    )
152    return torch_em.get_data_loader(dataset=dataset, batch_size=batch_size, **loader_kwargs)