torch_em.data.datasets.electron_microscopy.tnbc_mito

  1"""The TNBC-Mito dataset contains annotations for semantic segmentation of mitochondria in transmission
  2electron microscopy (TEM) images of triple-negative breast cancer (TNBC) models and Drp1-deficient mouse
  3skeletal muscle.
  4
  5The data covers four cohorts: DRP1-KO (Drp1-deficient mouse primary skeletal muscle cells), HCI-010 and
  6PIM001-P (TNBC patient-derived xenograft models), and Mixture (a pool of TNBC cell-line and xenograft images).
  7
  8The dataset is available at https://www.ebi.ac.uk/biostudies/bioimages/studies/S-BIAD2271 under the CC0 license.
  9The dataset was published in https://doi.org/10.1101/2025.02.19.635300.
 10Please cite this publication if you use the dataset in your research.
 11"""
 12
 13import os
 14from glob import glob
 15from typing import List, Literal, Optional, Tuple, Union
 16
 17from torch.utils.data import DataLoader, Dataset
 18
 19import torch_em
 20
 21from .. import util
 22
 23
 24BASE_URL = "https://ftp.ebi.ac.uk/biostudies/fire/S-BIAD/271/S-BIAD2271/Files"
 25MANIFEST_URL = f"{BASE_URL}/tem-seg-data_mitochondria_masks.tsv"
 26
 27COHORTS = ("DRP1-KO", "HCI-010", "Mixture", "PIM001-P")
 28SPLIT_DIRS = {"tra_val": "tra_val", "tst": "tst"}
 29
 30
 31def _read_manifest(manifest_path):
 32    """Parse the mask-to-source-image manifest into (mask_path, image_path) tuples."""
 33    pairs = []
 34    with open(manifest_path) as f:
 35        next(f)  # Skip the header line.
 36        for line in f:
 37            mask_path, image_path = line.strip().split("\t")
 38            pairs.append((mask_path, image_path))
 39    return pairs
 40
 41
 42def get_tnbc_mito_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 43    """Download the TNBC-Mito dataset.
 44
 45    Args:
 46        path: Filepath to a folder where the downloaded data will be saved.
 47        download: Whether to download the data if it is not present.
 48
 49    Returns:
 50        The filepath to the folder that mirrors the remote 'tem-seg/data' directory structure.
 51    """
 52    data_root = os.path.join(path, "tem-seg", "data")
 53    marker = os.path.join(path, ".download_complete")
 54    if os.path.exists(marker):
 55        return data_root
 56
 57    os.makedirs(path, exist_ok=True)
 58
 59    manifest_path = os.path.join(path, "tem-seg-data_mitochondria_masks.tsv")
 60    util.download_source(manifest_path, MANIFEST_URL, download, checksum=None)
 61    pairs = _read_manifest(manifest_path)
 62
 63    for mask_rel, image_rel in pairs:
 64        mask_path = os.path.join(path, mask_rel)
 65        image_path = os.path.join(path, image_rel)
 66        os.makedirs(os.path.dirname(mask_path), exist_ok=True)
 67        os.makedirs(os.path.dirname(image_path), exist_ok=True)
 68        util.download_source(mask_path, f"{BASE_URL}/{mask_rel}", download, checksum=None)
 69        util.download_source(image_path, f"{BASE_URL}/{image_rel}", download, checksum=None)
 70        _normalize_mask(mask_path)
 71
 72    with open(marker, "w"):
 73        pass
 74
 75    return data_root
 76
 77
 78def _normalize_mask(mask_path):
 79    """Re-save a palette-indexed mask PNG as a single-channel TIFF.
 80
 81    The source masks are palette ('P' mode) PNGs. Generic image readers expand palette images to RGB,
 82    which breaks single-channel label loading, so this writes the raw index values (already 0/1) once
 83    as a plain grayscale TIFF next to the original file.
 84    """
 85    import numpy as np
 86    import tifffile
 87    from PIL import Image
 88
 89    normalized_path = f"{os.path.splitext(mask_path)[0]}.tif"
 90    if os.path.exists(normalized_path):
 91        return
 92    labels = np.array(Image.open(mask_path))
 93    tifffile.imwrite(normalized_path, labels.astype("uint8"))
 94
 95
 96def get_tnbc_mito_paths(
 97    path: Union[os.PathLike, str],
 98    split: Literal["tra_val", "tst"],
 99    cohort: Optional[str] = None,
100    download: bool = False,
101) -> Tuple[List[str], List[str]]:
102    """Get paths to the TNBC-Mito raw images and mitochondria masks.
103
104    Args:
105        path: Filepath to a folder where the downloaded data will be saved.
106        split: The data split. Either 'tra_val' (pooled train / validation) or 'tst' (test).
107        cohort: The cohort to restrict to. One of 'DRP1-KO', 'HCI-010', 'Mixture', 'PIM001-P'.
108            If None, uses all cohorts that provide the requested split.
109        download: Whether to download the data if it is not present.
110
111    Returns:
112        The list of raw image paths.
113        The list of mitochondria mask paths.
114    """
115    assert split in SPLIT_DIRS, f"split must be one of {list(SPLIT_DIRS)}, got {split!r}"
116    if cohort is not None:
117        assert cohort in COHORTS, f"cohort must be one of {COHORTS}, got {cohort!r}"
118
119    data_root = get_tnbc_mito_data(path, download)
120
121    cohorts = [cohort] if cohort is not None else list(COHORTS)
122    raw_paths, label_paths = [], []
123    for c in cohorts:
124        image_dir = os.path.join(data_root, c, split, "slide_images")
125        mask_dir = os.path.join(data_root, c, split, "mitochondria", "masks")
126        if not os.path.exists(image_dir):
127            continue
128        images = sorted(glob(os.path.join(image_dir, "*.tif")))
129        for image_path in images:
130            fname = os.path.splitext(os.path.basename(image_path))[0]
131            mask_path = os.path.join(mask_dir, f"{fname}.tif")
132            assert os.path.exists(mask_path), f"Missing normalized mask for '{image_path}' at '{mask_path}'"
133            raw_paths.append(image_path)
134            label_paths.append(mask_path)
135
136    assert len(raw_paths) > 0, f"No images found for split '{split}' (cohort={cohort}) in '{data_root}'"
137    return raw_paths, label_paths
138
139
140def get_tnbc_mito_dataset(
141    path: Union[os.PathLike, str],
142    split: Literal["tra_val", "tst"],
143    patch_shape: Tuple[int, int],
144    cohort: Optional[str] = None,
145    download: bool = False,
146    **kwargs,
147) -> Dataset:
148    """Get the TNBC-Mito dataset for mitochondria segmentation in TEM images.
149
150    Args:
151        path: Filepath to a folder where the downloaded data will be saved.
152        split: The data split. Either 'tra_val' (pooled train / validation) or 'tst' (test).
153        patch_shape: The patch shape to use for training.
154        cohort: The cohort to restrict to. One of 'DRP1-KO', 'HCI-010', 'Mixture', 'PIM001-P'.
155            If None, uses all cohorts that provide the requested split.
156        download: Whether to download the data if it is not present.
157        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
158
159    Returns:
160        The segmentation dataset.
161    """
162    raw_paths, label_paths = get_tnbc_mito_paths(path, split, cohort, download)
163    return torch_em.default_segmentation_dataset(
164        raw_paths=raw_paths,
165        raw_key=None,
166        label_paths=label_paths,
167        label_key=None,
168        patch_shape=patch_shape,
169        is_seg_dataset=False,
170        **kwargs,
171    )
172
173
174def get_tnbc_mito_loader(
175    path: Union[os.PathLike, str],
176    split: Literal["tra_val", "tst"],
177    patch_shape: Tuple[int, int],
178    batch_size: int,
179    cohort: Optional[str] = None,
180    download: bool = False,
181    **kwargs,
182) -> DataLoader:
183    """Get the TNBC-Mito dataloader for mitochondria segmentation in TEM images.
184
185    Args:
186        path: Filepath to a folder where the downloaded data will be saved.
187        split: The data split. Either 'tra_val' (pooled train / validation) or 'tst' (test).
188        patch_shape: The patch shape to use for training.
189        batch_size: The batch size for training.
190        cohort: The cohort to restrict to. One of 'DRP1-KO', 'HCI-010', 'Mixture', 'PIM001-P'.
191            If None, uses all cohorts that provide the requested split.
192        download: Whether to download the data if it is not present.
193        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch
194            DataLoader.
195
196    Returns:
197        The PyTorch DataLoader.
198    """
199    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
200    dataset = get_tnbc_mito_dataset(path, split, patch_shape, cohort=cohort, download=download, **ds_kwargs)
201    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)