torch_em.data.datasets.light_microscopy.organoidnet

  1"""The OrganoIDNet dataset contains annotations of panceratic organoids.
  2
  3This dataset is from the publication https://doi.org/10.1007/s13402-024-00958-2.
  4Please cite it if you use this dataset for a publication.
  5"""
  6
  7
  8import os
  9import shutil
 10import zipfile
 11from glob import glob
 12from typing import Tuple, Union, List
 13
 14from torch.utils.data import Dataset, DataLoader
 15
 16import torch_em
 17
 18from .. import util
 19
 20
 21URL = "https://zenodo.org/records/10643410/files/OrganoIDNetData.zip?download=1"
 22CHECKSUM = "3cd9239bf74bda096ecb5b7bdb95f800c7fa30b9937f9aba6ddf98d754cbfa3d"
 23
 24
 25def get_organoidnet_data(path: Union[os.PathLike, str], split: str, download: bool = False) -> str:
 26    """Download the OrganoIDNet dataset.
 27
 28    Args:
 29        path: Filepath to the folder where the downloaded data will be saved.
 30        split: The data split to use.
 31        download: Whether to download the data if it is not present.
 32
 33    Returns:
 34        The filepath where the data is downloaded.
 35    """
 36    splits = ["Training", "Validation", "Test"]
 37    assert split in splits
 38
 39    os.makedirs(path, exist_ok=True)
 40
 41    data_dir = os.path.join(path, split)
 42    if os.path.exists(data_dir):
 43        return data_dir
 44
 45    # Download and extraction.
 46    zip_path = os.path.join(path, "OrganoIDNetData.zip")
 47    util.download_source(path=zip_path, url=URL, download=download, checksum=CHECKSUM)
 48
 49    # Only "Training", "Test", "Validation" from the zip are relevant and need to be extracted.
 50    # They are in "/OrganoIDNetData/Dataset/"
 51    prefix = "OrganoIDNetData/Dataset/"
 52    for dl_split in splits:
 53
 54        dl_prefix = prefix + dl_split
 55
 56        with zipfile.ZipFile(zip_path) as archive:
 57            for ff in archive.namelist():
 58                if ff.startswith(dl_prefix):
 59                    archive.extract(ff, path)
 60
 61    for dl_split in splits:
 62        shutil.move(
 63            os.path.join(path, "OrganoIDNetData/Dataset", dl_split),
 64            os.path.join(path, dl_split)
 65        )
 66
 67    assert os.path.exists(data_dir)
 68
 69    # os.remove(zip_path)
 70    return data_dir
 71
 72
 73def get_organoidnet_paths(
 74    path: Union[os.PathLike, str], split: str, download: bool = False
 75) -> Tuple[List[str], List[str]]:
 76    """Get paths to the OrganoIDNet data.
 77
 78    Args:
 79        path: Filepath to the folder where the downloaded data will be saved.
 80        split: The data split to use.
 81        download: Whether to download the data if it is not present.
 82
 83    Returns:
 84        List of filepaths for the image data.
 85        List of filepaths for the label data.
 86    """
 87    data_dir = get_organoidnet_data(path=path, split=split, download=download)
 88
 89    image_paths = sorted(glob(os.path.join(data_dir, "Images", "*.tif")))
 90    label_paths = sorted(glob(os.path.join(data_dir, "Masks", "*.tif")))
 91
 92    return image_paths, label_paths
 93
 94
 95def get_organoidnet_dataset(
 96    path: Union[os.PathLike, str], split: str, patch_shape: Tuple[int, int], download: bool = False, **kwargs
 97) -> Dataset:
 98    """Get the OrganoIDNet dataset for organoid segmentation in microscopy images.
 99
100    Args:
101        path: Filepath to a folder where the downloaded data will be saved.
102        split: The data split to use.
103        patch_shape: The patch shape to use for training.
104        download: Whether to download the data if it is not present.
105        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
106
107    Returns:
108        The segmentation dataset.
109    """
110    image_paths, label_paths = get_organoidnet_paths(path, split, download)
111
112    return torch_em.default_segmentation_dataset(
113        raw_paths=image_paths,
114        raw_key=None,
115        label_paths=label_paths,
116        label_key=None,
117        patch_shape=patch_shape,
118        is_seg_dataset=False,
119        **kwargs
120    )
121
122
123def get_organoidnet_loader(
124    path: Union[os.PathLike, str],
125    split: str,
126    patch_shape: Tuple[int, int],
127    batch_size: int,
128    download: bool = False,
129    **kwargs
130) -> DataLoader:
131    """Get the OrganoIDNet dataset for organoid segmentation in microscopy images.
132
133    Args:
134        path: Filepath to a folder where the downloaded data will be saved.
135        split: The data split to use.
136        patch_shape: The patch shape to use for training.
137        batch_size: The batch size for training.
138        download: Whether to download the data if it is not present.
139        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
140
141    Returns:
142        The DataLoader.
143    """
144    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
145    dataset = get_organoidnet_dataset(
146        path=path, split=split, patch_shape=patch_shape, download=download, **ds_kwargs
147    )
148    return torch_em.get_data_loader(dataset=dataset, batch_size=batch_size, **loader_kwargs)