torch_em.data.datasets.electron_microscopy.plantorganelle_hunter

  1"""The Plantorganelle Hunter dataset contains annotations for semantic segmentation of chloroplasts, mitochondria,
  2vacuoles and nuclei in transmission electron microscopy (TEM) images of plant cells.
  3
  4There are 902 large 2D micrographs (mostly 4800x3230 pixels), split into 541 train, 180 val and 181 test images.
  5
  6The dataset is from the publication "Plantorganelle Hunter is an effective deep-learning-based method for plant
  7organelle phenotyping in electron microscopy" and is available at https://doi.org/10.57760/sciencedb.01335.
  8Please cite it if you use this dataset in your research.
  9"""
 10
 11import os
 12from functools import partial
 13from typing import List, Literal, Optional, Tuple, Union
 14
 15import imageio.v3 as imageio
 16import numpy as np
 17from PIL import Image
 18from torch.utils.data import Dataset, DataLoader
 19
 20import torch_em
 21
 22from .. import util
 23
 24
 25# NOTE: The data is hosted on ScienceDB, which does not provide stable direct download links.
 26# The full archive is ~50 GB, but most of it is not used here (3D stacks, unseen data, ...).
 27DOI = "https://doi.org/10.57760/sciencedb.01335"
 28
 29# The mask files are palette images, with one (VOC-style) color per class.
 30CLASSES = {"chloroplast": 1, "mitochondria": 2, "vacuole": 3, "nucleus": 4}
 31COLORS_TO_CLASSES = {(128, 0, 0): 1, (0, 128, 0): 2, (128, 128, 0): 3, (0, 0, 128): 4}
 32
 33
 34def _convert_label(label_path, out_path):
 35    # We have to convert the labels, because the palette images get loaded as multi-channel RGB images.
 36    mask = Image.open(label_path)
 37    palette = np.array(mask.getpalette(), dtype="uint8").reshape(-1, 3)
 38    mask = np.array(mask)
 39
 40    lut = np.zeros(256, dtype="uint8")
 41    for idx in np.unique(mask):
 42        color = tuple(int(c) for c in palette[idx])
 43        if color == (0, 0, 0):
 44            continue
 45        if color not in COLORS_TO_CLASSES:
 46            raise ValueError(f"Unexpected label color {color} in {label_path}.")
 47        lut[idx] = COLORS_TO_CLASSES[color]
 48
 49    imageio.imwrite(out_path, lut[mask], compression=9)
 50
 51
 52def get_plantorganelle_hunter_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 53    """Prepare the Plantorganelle Hunter dataset.
 54
 55    The data cannot be downloaded automatically. Please download the archive from
 56    https://doi.org/10.57760/sciencedb.01335 and extract the folders 'image', 'label' and 'splits' to `path`.
 57
 58    Args:
 59        path: Filepath to the folder with the extracted data.
 60        download: Whether to download the data if it is not present. This is not supported for this dataset.
 61
 62    Returns:
 63        The filepath to the folder with the converted label masks.
 64    """
 65    label_dir = os.path.join(path, "label_classes")
 66    if os.path.exists(os.path.join(label_dir, ".done")):
 67        return label_dir
 68
 69    missing = [name for name in ("image", "label", "splits") if not os.path.exists(os.path.join(path, name))]
 70    if missing:
 71        raise RuntimeError(
 72            "The Plantorganelle Hunter dataset cannot be downloaded automatically. "
 73            f"Please download it from {DOI} and extract the folders 'image', 'label' and 'splits' to '{path}'. "
 74            f"Currently missing: {missing}."
 75        )
 76
 77    os.makedirs(label_dir, exist_ok=True)
 78    for split in ("train", "val", "test"):
 79        for name in _read_split(path, split):
 80            out_path = os.path.join(label_dir, f"{name}.png")
 81            if not os.path.exists(out_path):
 82                _convert_label(os.path.join(path, "label", f"{name}.png"), out_path)
 83
 84    open(os.path.join(label_dir, ".done"), "w").close()
 85    return label_dir
 86
 87
 88def _read_split(path, split):
 89    with open(os.path.join(path, "splits", f"{split}.txt")) as f:
 90        return [line.strip() for line in f if line.strip()]
 91
 92
 93def get_plantorganelle_hunter_paths(
 94    path: Union[os.PathLike, str],
 95    split: Literal["train", "val", "test"],
 96    download: bool = False,
 97) -> Tuple[List[str], List[str]]:
 98    """Get paths to the Plantorganelle Hunter data.
 99
100    Args:
101        path: Filepath to the folder with the extracted data.
102        split: The data split. Either 'train', 'val' or 'test'.
103        download: Whether to download the data if it is not present. This is not supported for this dataset.
104
105    Returns:
106        List of filepaths for the image data.
107        List of filepaths for the label data.
108    """
109    if split not in ("train", "val", "test"):
110        raise ValueError(f"'{split}' is not a valid split. Please choose from 'train' / 'val' / 'test'.")
111
112    label_dir = get_plantorganelle_hunter_data(path, download)
113
114    names = _read_split(path, split)
115    raw_paths = [os.path.join(path, "image", f"{name}.tif") for name in names]
116    label_paths = [os.path.join(label_dir, f"{name}.png") for name in names]
117    assert raw_paths and all(os.path.exists(p) for p in raw_paths), "Some of the images in the split are missing."
118
119    return raw_paths, label_paths
120
121
122def _select_class(labels, class_id):
123    return (labels == class_id).astype("uint8")
124
125
126def get_plantorganelle_hunter_dataset(
127    path: Union[os.PathLike, str],
128    patch_shape: Tuple[int, int],
129    split: Literal["train", "val", "test"],
130    label_choice: Optional[str] = None,
131    download: bool = False,
132    **kwargs
133) -> Dataset:
134    """Get the Plantorganelle Hunter dataset for organelle segmentation in plant cells.
135
136    Args:
137        path: Filepath to the folder with the extracted data.
138        patch_shape: The patch shape to use for training.
139        split: The data split. Either 'train', 'val' or 'test'.
140        label_choice: The structure to segment. One of 'chloroplast', 'mitochondria', 'vacuole' or 'nucleus'.
141            By default (None), all structures are returned as a semantic segmentation
142            (0: background, 1: chloroplast, 2: mitochondria, 3: vacuole, 4: nucleus).
143        download: Whether to download the data if it is not present. This is not supported for this dataset.
144        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
145
146    Returns:
147        The segmentation dataset.
148    """
149    if label_choice is not None:
150        if label_choice not in CLASSES:
151            raise ValueError(f"'{label_choice}' is not a valid label choice. Please choose from {list(CLASSES)}.")
152        kwargs = util.update_kwargs(kwargs, "label_transform", partial(_select_class, class_id=CLASSES[label_choice]))
153
154    raw_paths, label_paths = get_plantorganelle_hunter_paths(path, split, download)
155
156    return torch_em.default_segmentation_dataset(
157        raw_paths=raw_paths,
158        raw_key=None,
159        label_paths=label_paths,
160        label_key=None,
161        patch_shape=patch_shape,
162        is_seg_dataset=False,
163        **kwargs
164    )
165
166
167def get_plantorganelle_hunter_loader(
168    path: Union[os.PathLike, str],
169    batch_size: int,
170    patch_shape: Tuple[int, int],
171    split: Literal["train", "val", "test"],
172    label_choice: Optional[str] = None,
173    download: bool = False,
174    **kwargs
175) -> DataLoader:
176    """Get the Plantorganelle Hunter dataloader for organelle segmentation in plant cells.
177
178    Args:
179        path: Filepath to the folder with the extracted data.
180        batch_size: The batch size for training.
181        patch_shape: The patch shape to use for training.
182        split: The data split. Either 'train', 'val' or 'test'.
183        label_choice: The structure to segment. One of 'chloroplast', 'mitochondria', 'vacuole' or 'nucleus'.
184            By default (None), all structures are returned as a semantic segmentation
185            (0: background, 1: chloroplast, 2: mitochondria, 3: vacuole, 4: nucleus).
186        download: Whether to download the data if it is not present. This is not supported for this dataset.
187        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
188
189    Returns:
190        The DataLoader.
191    """
192    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
193    dataset = get_plantorganelle_hunter_dataset(path, patch_shape, split, label_choice, download, **ds_kwargs)
194    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)