torch_em.data.datasets.light_microscopy.orgasegment

  1"""The OrgaSegment dataset contains annotations for organoid segmentation
  2of intestinal patient derived organoids in bright field images.
  3
  4This dataset is from the publication https://doi.org/10.1038/s42003-024-05966-4.
  5Please cite it if you use this dataset for your research.
  6"""
  7
  8import os
  9import shutil
 10from glob import glob
 11from typing import Tuple, Union, Literal, List
 12
 13from torch.utils.data import Dataset, DataLoader
 14
 15import torch_em
 16
 17from .. import util
 18
 19
 20URL = "https://zenodo.org/records/10278229/files/OrganoidBasic_v20211206.zip"
 21CHECKSUM = "d067124d734108e46e18f65daaf17c89cb0a40bdacc6f6031815a6839e472798"
 22
 23
 24def get_orgasegment_data(
 25    path: Union[os.PathLike, str], split: Literal["train", "val", "eval"], download: bool = False
 26) -> str:
 27    """Download the OrgaSegment dataset for organoid segmentation.
 28
 29    Args:
 30        path: Filepath to a folder where the downloaded data will be saved.
 31        split: The split to download. Either 'train', 'val or 'eval'.
 32        download: Whether to download the data if it is not present.
 33
 34    Returns:
 35        The filepath to the training data.
 36    """
 37    os.makedirs(path, exist_ok=True)
 38
 39    data_dir = os.path.join(path, split)
 40    if os.path.exists(data_dir):
 41        return data_dir
 42
 43    zip_path = os.path.join(path, "OrganoidBasic_v20211206.zip")
 44    util.download_source(path=zip_path, url=URL, download=download, checksum=CHECKSUM)
 45    util.unzip(zip_path=zip_path, dst=path, remove=True)
 46
 47    shutil.move(os.path.join(path, "OrganoidBasic_v20211206", "train"), os.path.join(path, "train"))
 48    shutil.move(os.path.join(path, "OrganoidBasic_v20211206", "val"), os.path.join(path, "val"))
 49    shutil.move(os.path.join(path, "OrganoidBasic_v20211206", "eval"), os.path.join(path, "eval"))
 50    shutil.rmtree(os.path.join(path, "OrganoidBasic_v20211206"))
 51
 52    return data_dir
 53
 54
 55def get_orgasegment_paths(
 56    path: Union[os.PathLike, str], split: Literal["train", "val", "eval"], download: bool = False
 57) -> Tuple[List[str], List[str]]:
 58    """Get paths for the OrgaSegment data.
 59
 60    Args:
 61        path: Filepath to a folder where the downloaded data will be saved.
 62        split: The split to download. Either 'train', 'val or 'eval'.
 63        download: Whether to download the data if it is not present.
 64
 65    Returns:
 66        List of filepaths to the image data.
 67        List of filepaths to the label data.
 68    """
 69    data_dir = get_orgasegment_data(path=path, split=split, download=download)
 70
 71    image_paths = sorted(glob(os.path.join(data_dir, "*_img.jpg")))
 72    label_paths = sorted(glob(os.path.join(data_dir, "*_masks_organoid.png")))
 73
 74    return image_paths, label_paths
 75
 76
 77def get_orgasegment_dataset(
 78    path: Union[os.PathLike, str],
 79    patch_shape: Tuple[int, int],
 80    split: Literal["train", "val", "eval"],
 81    boundaries: bool = False,
 82    binary: bool = False,
 83    download: bool = False,
 84    **kwargs
 85) -> Dataset:
 86    """Get the OrgaSegment dataset for organoid segmentation
 87
 88    Args:
 89        path: Filepath to a folder where the downloaded data will be saved.
 90        patch_shape: The patch shape to use for training.
 91        split: The split to download. Either 'train', 'val or 'eval'.
 92        boundaries: Whether to compute boundaries as the target.
 93        binary: Whether to use a binary segmentation target.
 94        download: Whether to download the data if it is not present.
 95        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
 96
 97    Returns:
 98        The segmentation dataset.
 99    """
100    assert split in ["train", "val", "eval"]
101
102    image_paths, label_paths = get_orgasegment_paths(path, split, download)
103
104    kwargs, _ = util.add_instance_label_transform(kwargs, add_binary_target=True, binary=binary, boundaries=boundaries)
105
106    return torch_em.default_segmentation_dataset(
107        raw_paths=image_paths,
108        raw_key=None,
109        label_paths=label_paths,
110        label_key=None,
111        patch_shape=patch_shape,
112        is_seg_dataset=False,
113        **kwargs
114    )
115
116
117def get_orgasegment_loader(
118    path: Union[os.PathLike, str],
119    batch_size: int,
120    patch_shape: Tuple[int, int],
121    split: Literal["train", "val", "eval"],
122    boundaries: bool = False,
123    binary: bool = False,
124    download: bool = False,
125    **kwargs
126) -> DataLoader:
127    """Get the OrgaSegment dataloader for organoid segmentation
128
129    Args:
130        path: Filepath to a folder where the downloaded data will be saved.
131        batch_size: The batch size for training.
132        patch_shape: The patch shape to use for training.
133        split: The split to download. Either 'train', 'val or 'eval'.
134        boundaries: Whether to compute boundaries as the target.
135        binary: Whether to use a binary segmentation target.
136        download: Whether to download the data if it is not present.
137        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
138
139    Returns:
140        The DataLoader.
141    """
142    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
143    dataset = get_orgasegment_dataset(path, patch_shape, split, boundaries, binary, download, **ds_kwargs)
144    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)