torch_em.data.datasets.electron_microscopy.hela_mito

  1"""The HeLa-Mito dataset contains mitochondria instance segmentation for 2D electron
  2microscopy slices of HeLa cells.
  3
  4The raw slices are from the same acquisition (EMPIAR-10094) already used by the
  5`cefa_hela` module, but that module only provides nuclear-envelope annotations; this
  6module adds a separate, complementary layer of mitochondria instance masks on 9 of the
  724 raw slices released alongside it.
  8
  9Annotation coverage is sparse, not exhaustive: only 9 of the 24 available slices carry
 10labels, and some mitochondria visible in the raw data may be left unannotated even on
 11labeled slices.
 12
 13The dataset is available at https://github.com/reyesaldasoro/MitoEM.
 14The dataset was published in https://doi.org/10.1101/2023.11.14.567016.
 15Please cite this publication if you use the dataset in your research.
 16"""
 17
 18import os
 19from typing import List, Tuple, Union
 20
 21from torch.utils.data import DataLoader, Dataset
 22
 23import torch_em
 24
 25from .. import util
 26
 27
 28BASE_URL = "https://raw.githubusercontent.com/reyesaldasoro/MitoEM/main/CODE"
 29RAW_URL = BASE_URL + "/OriginalImages/ROI_6005_4739_81_z{slice_id}.tif"
 30LABEL_URL = BASE_URL + "/{prefix}_ROI_6005_4739_81_z{slice_id}.mat"
 31
 32# The annotated slice ids and the ground-truth file prefix to use for each. GT and GT2
 33# both annotate z0001 (0.94 IoU between them); GT is used there and GT2 only for the
 34# 4 additional slices it uniquely covers.
 35SLICES = {
 36    "0001": "GT",
 37    "0002": "GT2",
 38    "0004": "GT2",
 39    "0026": "GT2",
 40    "0028": "GT2",
 41    "0030": "GT",
 42    "0060": "GT",
 43    "0116": "GT",
 44    "0150": "GT",
 45}
 46
 47
 48def _convert_slice(raw_path, mat_path, out_path):
 49    import h5py
 50    import tifffile
 51    import scipy.io as sio
 52
 53    raw = tifffile.imread(raw_path)
 54    labels = sio.loadmat(mat_path)["groundTruthM"].max(axis=2)
 55    assert raw.shape == labels.shape, f"{raw.shape} != {labels.shape}"
 56
 57    tmp_path = out_path + ".incomplete"
 58    with h5py.File(tmp_path, "w") as f:
 59        f.create_dataset("raw", data=raw, compression="gzip")
 60        f.create_dataset("labels", data=labels, compression="gzip")
 61    os.rename(tmp_path, out_path)
 62
 63
 64def get_hela_mito_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 65    """Download the HeLa-Mito dataset.
 66
 67    Args:
 68        path: Filepath to a folder where the data will be downloaded.
 69        download: Whether to download the data if it is not present.
 70
 71    Returns:
 72        Filepath to the folder with the per-slice HDF5 files.
 73    """
 74    os.makedirs(path, exist_ok=True)
 75
 76    for slice_id, prefix in SLICES.items():
 77        out_path = os.path.join(path, f"z{slice_id}.h5")
 78        if os.path.exists(out_path):
 79            continue
 80        if not download:
 81            raise RuntimeError(f"Cannot find the data at {out_path}, but download was set to False.")
 82
 83        raw_path = os.path.join(path, f"raw_z{slice_id}.tif")
 84        mat_path = os.path.join(path, f"{prefix}_z{slice_id}.mat")
 85        util.download_source(raw_path, RAW_URL.format(slice_id=slice_id), download)
 86        util.download_source(mat_path, LABEL_URL.format(prefix=prefix, slice_id=slice_id), download)
 87
 88        _convert_slice(raw_path, mat_path, out_path)
 89        os.remove(raw_path)
 90        os.remove(mat_path)
 91
 92    return path
 93
 94
 95def get_hela_mito_paths(path: Union[os.PathLike, str], download: bool = False) -> List[str]:
 96    """Get paths to the HeLa-Mito data.
 97
 98    Args:
 99        path: Filepath to a folder where the data will be downloaded.
100        download: Whether to download the data if it is not present.
101
102    Returns:
103        List of filepaths to the per-slice HDF5 files.
104    """
105    data_dir = get_hela_mito_data(path, download)
106    return [os.path.join(data_dir, f"z{slice_id}.h5") for slice_id in SLICES]
107
108
109def get_hela_mito_dataset(
110    path: Union[os.PathLike, str],
111    patch_shape: Tuple[int, int],
112    download: bool = False,
113    **kwargs
114) -> Dataset:
115    """Get the dataset for mitochondria instance segmentation in HeLa cell EM slices.
116
117    Args:
118        path: Filepath to a folder where the data will be downloaded.
119        patch_shape: The patch shape to use for training.
120        download: Whether to download the data if it is not present.
121        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
122
123    Returns:
124        The segmentation dataset.
125    """
126    assert len(patch_shape) == 2
127
128    paths = get_hela_mito_paths(path, download)
129
130    kwargs = util.update_kwargs(kwargs, "is_seg_dataset", True)
131    kwargs, _ = util.add_instance_label_transform(kwargs, add_binary_target=True)
132
133    return torch_em.default_segmentation_dataset(
134        raw_paths=paths,
135        raw_key="raw",
136        label_paths=paths,
137        label_key="labels",
138        patch_shape=patch_shape,
139        **kwargs
140    )
141
142
143def get_hela_mito_loader(
144    path: Union[os.PathLike, str],
145    patch_shape: Tuple[int, int],
146    batch_size: int,
147    download: bool = False,
148    **kwargs
149) -> DataLoader:
150    """Get the DataLoader for mitochondria instance segmentation in HeLa cell EM slices.
151
152    Args:
153        path: Filepath to a folder where the data will be downloaded.
154        patch_shape: The patch shape to use for training.
155        batch_size: The batch size for training.
156        download: Whether to download the data if it is not present.
157        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
158
159    Returns:
160        The DataLoader.
161    """
162    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
163    dataset = get_hela_mito_dataset(path, patch_shape, download=download, **ds_kwargs)
164    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)