torch_em.data.datasets.electron_microscopy.cardiomyocyte_tem

  1"""The Cardiomyocyte TEM dataset contains annotations for semantic segmentation of cells, myofibres,
  2mitochondria, nuclei and nucleoli in transmission electron microscopy (TEM) images of isolated cardiomyocytes.
  3
  4There are nine large 2D micrographs (3500x magnification) with binary masks. Nucleoli are only annotated in one image.
  5
  6The dataset is from the publication "Comprehensive Transmission Electron Microscopy Platform for Spatially Resolved
  7Mitochondrial and Nuclear Architecture in Isolated Cardiomyocytes" and is available at
  8https://doi.org/10.5281/zenodo.17068504. Please cite it if you use this dataset in your research.
  9"""
 10
 11import os
 12from glob import glob
 13from typing import List, Tuple, Union
 14
 15from natsort import natsorted
 16from torch.utils.data import Dataset, DataLoader
 17
 18import torch_em
 19from torch_em.transform.label import labels_to_binary
 20
 21from .. import util
 22
 23
 24BASE_URL = "https://zenodo.org/records/17068504/files"
 25RAW_URLS = {
 26    "914.11_grid1_nucA_m2_M3500x": "61bad379ba42bcc0c8429448dab346811aefb25ff76d07ddf9670a960b99b07c",
 27    "914.11_grid4_nucA_m2_M3500x": "ee40d1a849df3e1efea0d36d152d63779179c587a03dafc9a44cdcb687eb881e",
 28    "914.16_18.12.2023_nucD_m2_3500x": "e21fec9443e08f0e8bb49e264622f2f89227fbad4f69126eb076199f9606252e",
 29    "914.1_nucA_m3_3500x": "eec1368185ce12825281af4a7dbe080332336957b94965af336035418fc32d96",
 30    "914.1_nucC_m3_3500x": "b4f676d18881b7a451d9a5ecc5f8ff3cafd913105a1bb43dcdd25af07a0aea8f",
 31    "914.22_nucC_m3_3500x": "a3370477636f69d707f2a79388b5a725ac5a87594ff7b171ba3b76fdd4669a93",
 32    "914.22_nucE_m2_3500x": "0b2ae97f1b60f1a910958c1478a8f427cdc17d5ef8ab65d0b32e52e689f79d38",
 33    "914.4_nucA_m3_M3500x": "2b0be4a9f3161be1a96c7ff29b6c3eaa42610fae3ff3c71c68d8e862bb109197",
 34    "914.9_grid1_nucA_m3_M3500x": "679ffd9fa4eaa619e4fe199bc84b8667d5ca53ef1a9f1c90cf0bbb967fe6e7d6",
 35}
 36# NOTE: The mask archives contain more images than the released raw data. Only the ones with raw data are used.
 37# Some of the masks have a 'KO_' prefix in their name, which is missing in the name of the corresponding raw image.
 38MASK_ZIPS = {
 39    "cell": ("Cell_Masks", "dff94c42d1150ffd78657429ba5695c7a2ed91522f5f224bd4ad1366f07ad796"),
 40    "myo": ("Myo_Masks", "c7e63c6421d0a2f5ecc129d930ab40c7df4c5768f142102a19c86ee674d8342e"),
 41    "mito": ("Mito_Masks", "50d43cb088ca0e3650881ee10139f0889197d60584994e18c29c95513537eecd"),
 42    "nucleus": ("Nucleus_Masks", "62ab8a2f1ad72bb2e5766ab00cdd48a74bdb71b6327761527d909551c35857e1"),
 43    "nucleolus": ("Nucleolus_Masks", "c52455cab0c8c35a219ad8823904a98ca041a6c2725550df059afd82c864d1c1"),
 44}
 45LABELS = list(MASK_ZIPS.keys())
 46
 47
 48def get_cardiomyocyte_tem_data(
 49    path: Union[os.PathLike, str], label_choice: str = "mito", download: bool = False
 50) -> str:
 51    """Download the Cardiomyocyte TEM dataset.
 52
 53    Args:
 54        path: Filepath to a folder where the downloaded data will be saved.
 55        label_choice: The structure to segment. One of 'cell', 'myo', 'mito', 'nucleus' or 'nucleolus'.
 56        download: Whether to download the data if it is not present.
 57
 58    Returns:
 59        The filepath to the data folder.
 60    """
 61    if label_choice not in LABELS:
 62        raise ValueError(f"'{label_choice}' is not a valid label choice. Please choose from {LABELS}.")
 63
 64    image_dir = os.path.join(path, "images")
 65    os.makedirs(image_dir, exist_ok=True)
 66    for name, checksum in RAW_URLS.items():
 67        util.download_source(
 68            os.path.join(image_dir, f"{name}.tif"), f"{BASE_URL}/{name}.tif?download=1", download, checksum
 69        )
 70
 71    folder, checksum = MASK_ZIPS[label_choice]
 72    if not os.path.exists(os.path.join(path, folder)):
 73        zip_path = os.path.join(path, f"{folder}.zip")
 74        util.download_source(zip_path, f"{BASE_URL}/{folder}.zip?download=1", download, checksum)
 75        util.unzip(zip_path, path)
 76
 77    return path
 78
 79
 80def get_cardiomyocyte_tem_paths(
 81    path: Union[os.PathLike, str], label_choice: str = "mito", download: bool = False
 82) -> Tuple[List[str], List[str]]:
 83    """Get paths to the Cardiomyocyte TEM data.
 84
 85    Args:
 86        path: Filepath to a folder where the downloaded data will be saved.
 87        label_choice: The structure to segment. One of 'cell', 'myo', 'mito', 'nucleus' or 'nucleolus'.
 88        download: Whether to download the data if it is not present.
 89
 90    Returns:
 91        List of filepaths for the image data.
 92        List of filepaths for the label data.
 93    """
 94    get_cardiomyocyte_tem_data(path, label_choice, download)
 95
 96    folder, _ = MASK_ZIPS[label_choice]
 97    mask_paths = glob(os.path.join(path, folder, "*.tif"))
 98
 99    raw_paths, label_paths = [], []
100    for name in RAW_URLS:
101        names = (f"{name}_{label_choice}.tif", f"KO_{name}_{label_choice}.tif")
102        match = [p for p in mask_paths if os.path.basename(p) in names]
103        # Not all images have masks for every structure (e.g. nucleoli).
104        if not match:
105            continue
106        raw_paths.append(os.path.join(path, "images", f"{name}.tif"))
107        label_paths.append(match[0])
108
109    assert raw_paths and len(raw_paths) == len(label_paths)
110    return natsorted(raw_paths), natsorted(label_paths, key=lambda p: os.path.basename(p).replace("KO_", ""))
111
112
113def get_cardiomyocyte_tem_dataset(
114    path: Union[os.PathLike, str],
115    patch_shape: Tuple[int, int],
116    label_choice: str = "mito",
117    download: bool = False,
118    **kwargs
119) -> Dataset:
120    """Get the Cardiomyocyte TEM dataset for cell and organelle segmentation.
121
122    Args:
123        path: Filepath to a folder where the downloaded data will be saved.
124        patch_shape: The patch shape to use for training.
125        label_choice: The structure to segment. One of 'cell', 'myo', 'mito', 'nucleus' or 'nucleolus'.
126        download: Whether to download the data if it is not present.
127        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
128
129    Returns:
130        The segmentation dataset.
131    """
132    raw_paths, label_paths = get_cardiomyocyte_tem_paths(path, label_choice, download)
133
134    kwargs = util.update_kwargs(kwargs, "label_transform", labels_to_binary)
135
136    return torch_em.default_segmentation_dataset(
137        raw_paths=raw_paths,
138        raw_key=None,
139        label_paths=label_paths,
140        label_key=None,
141        patch_shape=patch_shape,
142        is_seg_dataset=False,
143        **kwargs
144    )
145
146
147def get_cardiomyocyte_tem_loader(
148    path: Union[os.PathLike, str],
149    batch_size: int,
150    patch_shape: Tuple[int, int],
151    label_choice: str = "mito",
152    download: bool = False,
153    **kwargs
154) -> DataLoader:
155    """Get the Cardiomyocyte TEM dataloader for cell and organelle segmentation.
156
157    Args:
158        path: Filepath to a folder where the downloaded data will be saved.
159        batch_size: The batch size for training.
160        patch_shape: The patch shape to use for training.
161        label_choice: The structure to segment. One of 'cell', 'myo', 'mito', 'nucleus' or 'nucleolus'.
162        download: Whether to download the data if it is not present.
163        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
164
165    Returns:
166        The DataLoader.
167    """
168    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
169    dataset = get_cardiomyocyte_tem_dataset(path, patch_shape, label_choice, download, **ds_kwargs)
170    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)