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 5 of the
  724 raw slices released alongside it.
  8
  9Annotation coverage is sparse, not exhaustive: only 5 of the 24 available slices carry
 10labels, only the central cell is annotated, and some mitochondria visible in the raw data
 11may be left unannotated even in the central cell.
 12
 13The repository contains two independent annotations of these 5 slices ('gt' and 'gt2'), which can be selected
 14via `annotation_version`. They are useful to study annotation disagreement.
 15
 16The dataset is available at https://github.com/reyesaldasoro/MitoEM.
 17The dataset was published in https://doi.org/10.1101/2023.11.14.567016.
 18Please cite this publication if you use the dataset in your research.
 19"""
 20
 21import os
 22from typing import List, Literal, Tuple, Union
 23
 24from torch.utils.data import DataLoader, Dataset
 25
 26import torch_em
 27
 28from .. import util
 29
 30
 31BASE_URL = "https://raw.githubusercontent.com/reyesaldasoro/MitoEM/main/CODE"
 32RAW_URL = BASE_URL + "/OriginalImages/ROI_6005_4739_81_z{slice_id}.tif"
 33LABEL_URLS = {
 34    "gt": BASE_URL + "/GT_ROI_6005_4739_81_z{label_id}.mat",
 35    "gt2": BASE_URL + "/GT2_ROI_6005_4739_81_z{label_id}.mat",
 36}
 37
 38# The annotated raw slice ids, with the 'GT' annotation of each. The repository also has 'GT2' files, but these are
 39# a second annotation of the same 5 slices and numbered sequentially, not by raw slice: GT2 z0001, z0002, z0004,
 40# z0026 and z0028 annotate the raw slices z0001, z0030, z0060, z0116 and z0150 (mask IoU with GT 0.76 to 0.96).
 41SLICES = ["0001", "0030", "0060", "0116", "0150"]
 42LABEL_IDS = {
 43    "gt": {slice_id: slice_id for slice_id in SLICES},
 44    "gt2": dict(zip(SLICES, ["0001", "0002", "0004", "0026", "0028"])),
 45}
 46
 47
 48def _check_annotation_version(annotation_version):
 49    if annotation_version not in LABEL_URLS:
 50        raise ValueError(f"Invalid annotation_version '{annotation_version}', choose one of {list(LABEL_URLS)}.")
 51
 52
 53def _slice_path(path, slice_id, annotation_version):
 54    # The 'gt' files keep their original name, so that existing downloads stay valid.
 55    suffix = "" if annotation_version == "gt" else f"_{annotation_version}"
 56    return os.path.join(path, f"z{slice_id}{suffix}.h5")
 57
 58
 59def _convert_slice(raw_path, mat_path, out_path):
 60    import h5py
 61    import tifffile
 62    import scipy.io as sio
 63
 64    raw = tifffile.imread(raw_path)
 65    labels = sio.loadmat(mat_path)["groundTruthM"].max(axis=2)
 66    assert raw.shape == labels.shape, f"{raw.shape} != {labels.shape}"
 67
 68    tmp_path = out_path + ".incomplete"
 69    with h5py.File(tmp_path, "w") as f:
 70        f.create_dataset("raw", data=raw, compression="gzip")
 71        f.create_dataset("labels", data=labels, compression="gzip")
 72    os.rename(tmp_path, out_path)
 73
 74
 75def get_hela_mito_data(
 76    path: Union[os.PathLike, str], download: bool = False, annotation_version: Literal["gt", "gt2"] = "gt"
 77) -> str:
 78    """Download the HeLa-Mito dataset.
 79
 80    Args:
 81        path: Filepath to a folder where the data will be downloaded.
 82        download: Whether to download the data if it is not present.
 83        annotation_version: The annotation set to use, either 'gt' or 'gt2' (a second, independent annotation
 84            of the same 5 slices).
 85
 86    Returns:
 87        Filepath to the folder with the per-slice HDF5 files.
 88    """
 89    _check_annotation_version(annotation_version)
 90    os.makedirs(path, exist_ok=True)
 91
 92    for slice_id in SLICES:
 93        out_path = _slice_path(path, slice_id, annotation_version)
 94        if os.path.exists(out_path):
 95            continue
 96        if not download:
 97            raise RuntimeError(f"Cannot find the data at {out_path}, but download was set to False.")
 98
 99        raw_path = os.path.join(path, f"raw_z{slice_id}.tif")
100        mat_path = os.path.join(path, f"{annotation_version}_z{slice_id}.mat")
101        label_url = LABEL_URLS[annotation_version].format(label_id=LABEL_IDS[annotation_version][slice_id])
102        util.download_source(raw_path, RAW_URL.format(slice_id=slice_id), download)
103        util.download_source(mat_path, label_url, download)
104
105        _convert_slice(raw_path, mat_path, out_path)
106        os.remove(raw_path)
107        os.remove(mat_path)
108
109    return path
110
111
112def get_hela_mito_paths(
113    path: Union[os.PathLike, str], download: bool = False, annotation_version: Literal["gt", "gt2"] = "gt"
114) -> List[str]:
115    """Get paths to the HeLa-Mito data.
116
117    Args:
118        path: Filepath to a folder where the data will be downloaded.
119        download: Whether to download the data if it is not present.
120        annotation_version: The annotation set to use, either 'gt' or 'gt2'.
121
122    Returns:
123        List of filepaths to the per-slice HDF5 files.
124    """
125    data_dir = get_hela_mito_data(path, download, annotation_version)
126    return [_slice_path(data_dir, slice_id, annotation_version) for slice_id in SLICES]
127
128
129def get_hela_mito_dataset(
130    path: Union[os.PathLike, str],
131    patch_shape: Tuple[int, int],
132    download: bool = False,
133    annotation_version: Literal["gt", "gt2"] = "gt",
134    **kwargs
135) -> Dataset:
136    """Get the dataset for mitochondria instance segmentation in HeLa cell EM slices.
137
138    Args:
139        path: Filepath to a folder where the data will be downloaded.
140        patch_shape: The patch shape to use for training.
141        download: Whether to download the data if it is not present.
142        annotation_version: The annotation set to use, either 'gt' or 'gt2'.
143        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
144
145    Returns:
146        The segmentation dataset.
147    """
148    assert len(patch_shape) == 2
149
150    paths = get_hela_mito_paths(path, download, annotation_version)
151
152    kwargs = util.update_kwargs(kwargs, "is_seg_dataset", True)
153    kwargs, _ = util.add_instance_label_transform(kwargs, add_binary_target=True)
154
155    return torch_em.default_segmentation_dataset(
156        raw_paths=paths,
157        raw_key="raw",
158        label_paths=paths,
159        label_key="labels",
160        patch_shape=patch_shape,
161        **kwargs
162    )
163
164
165def get_hela_mito_loader(
166    path: Union[os.PathLike, str],
167    patch_shape: Tuple[int, int],
168    batch_size: int,
169    download: bool = False,
170    annotation_version: Literal["gt", "gt2"] = "gt",
171    **kwargs
172) -> DataLoader:
173    """Get the DataLoader for mitochondria instance segmentation in HeLa cell EM slices.
174
175    Args:
176        path: Filepath to a folder where the data will be downloaded.
177        patch_shape: The patch shape to use for training.
178        batch_size: The batch size for training.
179        download: Whether to download the data if it is not present.
180        annotation_version: The annotation set to use, either 'gt' or 'gt2'.
181        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
182
183    Returns:
184        The DataLoader.
185    """
186    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
187    dataset = get_hela_mito_dataset(
188        path, patch_shape, download=download, annotation_version=annotation_version, **ds_kwargs
189    )
190    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)