torch_em.data.datasets.electron_microscopy.densecell

  1"""The DenseCell dataset contains annotations for semantic segmentation of densely-packed cellular organelles
  2in serial block-face scanning electron microscopy (SBF-SEM) images of platelet tissue.
  3
  4The dataset was published in https://doi.org/10.1038/s41598-021-81590-0.
  5Please cite this publication if you use the dataset in your research.
  6"""
  7
  8import os
  9from shutil import rmtree
 10from typing import Tuple, Union, Literal, Optional
 11
 12import numpy as np
 13from scipy import ndimage
 14from skimage.feature import peak_local_max
 15from skimage.segmentation import watershed
 16from skimage.measure import label as connected_components
 17
 18from torch.utils.data import Dataset, DataLoader
 19
 20import torch_em
 21
 22from .. import util
 23
 24
 25URL = "https://www.dropbox.com/s/68yclbraqq1diza/platelet_data_1219.zip?dl=1"
 26CHECKSUM = None
 27
 28ORGANELLES = {
 29    1: "cell",
 30    2: "mitochondrion",
 31    3: "alpha_granule",
 32    4: "canalicular_vessel",
 33    5: "dense_granule",
 34    6: "dense_core",
 35}
 36
 37SPLIT_FILES = {
 38    "train": {"images": "train-images.tif", "labels": "train-labels.tif"},
 39    "val": {"images": "eval-images.tif", "labels": "eval-labels.tif"},
 40    "test": {"images": "test-images.tif", "labels": "test-labels.tif"},
 41}
 42
 43
 44def get_densecell_data(
 45    path: Union[os.PathLike, str], split: Literal["train", "val", "test"], download: bool = False
 46) -> str:
 47    """Download the DenseCell dataset.
 48
 49    Args:
 50        path: Filepath to a folder where the downloaded data will be saved.
 51        split: The split to download. Either 'train', 'val', or 'test'.
 52        download: Whether to download the data if it is not present.
 53
 54    Returns:
 55        The filepath for the downloaded data.
 56    """
 57    import h5py
 58    import tifffile
 59
 60    data_path = os.path.join(path, f"densecell_{split}.h5")
 61    if os.path.exists(data_path):
 62        with h5py.File(data_path, "r") as f:
 63            if "labels/original" in f:
 64                return data_path
 65
 66        # Remove old file with outdated structure.
 67        os.remove(data_path)
 68
 69    os.makedirs(path, exist_ok=True)
 70
 71    # Download and extract the ZIP if the source TIFFs are not available.
 72    platelet_dir = os.path.join(path, "platelet_data")
 73    if not os.path.exists(platelet_dir):
 74        zip_path = os.path.join(path, "platelet_data_1219.zip")
 75        util.download_source(zip_path, URL, download, checksum=CHECKSUM)
 76        util.unzip(zip_path, path, remove=True)
 77
 78    assert os.path.exists(platelet_dir), f"Expected extracted directory at {platelet_dir}"
 79
 80    for _split, files in SPLIT_FILES.items():
 81        out_path = os.path.join(path, f"densecell_{_split}.h5")
 82        if os.path.exists(out_path):
 83            with h5py.File(out_path, "r") as f:
 84                if "labels/original" in f:
 85                    continue
 86
 87            os.remove(out_path)
 88
 89        raw = tifffile.imread(os.path.join(platelet_dir, files["images"]))
 90        labels = tifffile.imread(os.path.join(platelet_dir, files["labels"]))
 91        assert raw.shape == labels.shape, f"Shape mismatch for {_split}: {raw.shape} vs {labels.shape}"
 92
 93        labels = labels.astype(np.uint8)
 94        with h5py.File(out_path, "w") as f:
 95            f.create_dataset("raw", data=raw, compression="gzip")
 96            f.create_dataset("labels/original", data=labels, compression="gzip")
 97            for label_id, name in ORGANELLES.items():
 98                # For cells, use all non-background labels to avoid holes from internal organelles.
 99                if name == "cell":
100                    binary_mask = (labels >= 1).astype(np.uint8)
101                else:
102                    binary_mask = (labels == label_id).astype(np.uint8)
103
104                f.create_dataset(f"labels/{name}", data=binary_mask, compression="gzip")
105
106    rmtree(platelet_dir)
107
108    assert os.path.exists(data_path), data_path
109    return data_path
110
111
112CELL_INSTANCE_KEY = "labels/cell_instances"
113
114
115def _merge_without_neck(segments, distance, max_saddle_ratio, min_peak, min_z_overlap=0.5):
116    """Merge watershed segments of the same cell.
117
118    In-plane neighbours merge when the distance saddle on their boundary is high relative to their peaks (no
119    neck). Neighbours across z merge when their footprints overlap by at least *min_z_overlap* IoU. Segments
120    with a peak below *min_peak* are slivers that would bridge cells, so they are absorbed by their largest
121    contact instead.
122    """
123    peaks = ndimage.maximum(distance, segments, index=np.arange(segments.max() + 1))
124    small = peaks < min_peak
125    merged = np.arange(segments.max() + 1)
126
127    def union(u, v):
128        ru, rv = merged[u], merged[v]
129        if ru != rv:
130            merged[merged == max(ru, rv)] = min(ru, rv)
131
132    contact = {}
133    for axis in range(1, segments.ndim):
134        a = np.moveaxis(segments, axis, 0)
135        d = np.moveaxis(distance, axis, 0)
136        left, right = a[:-1], a[1:]
137        touch = (left != right) & (left > 0) & (right > 0)
138        saddle = np.minimum(d[:-1], d[1:])[touch]
139        for u, v, sd in zip(left[touch], right[touch], saddle):
140            if small[u] or small[v]:
141                contact[(u, v)] = contact.get((u, v), 0) + 1
142                contact[(v, u)] = contact.get((v, u), 0) + 1
143            elif sd > max_saddle_ratio * min(peaks[u], peaks[v]):
144                union(u, v)
145
146    for z in range(segments.shape[0] - 1):
147        lower, upper = segments[z], segments[z + 1]
148        touch = (lower != upper) & (lower > 0) & (upper > 0)
149        for u, v in set(zip(lower[touch].tolist(), upper[touch].tolist())):
150            if small[u] or small[v]:
151                continue
152            fu, fv = lower == u, upper == v
153            iou = np.logical_and(fu, fv).sum() / np.logical_or(fu, fv).sum()
154            if iou >= min_z_overlap:
155                union(u, v)
156
157    for u in np.nonzero(small)[0]:
158        neighbours = [(n, c) for (a, n), c in contact.items() if a == u and not small[n]]
159        if neighbours:
160            union(u, max(neighbours, key=lambda item: item[1])[0])
161    return merged[segments]
162
163
164def _derive_cell_instances(cell_mask, sampling=(5, 1, 1), min_distance=60, max_saddle_ratio=0.6, min_size=2000):
165    """Split the binary cell mask into instances with a seeded 3D watershed.
166
167    Seeds are per-section 2D distance maxima at least *min_distance* apart, so small cells next to a large
168    neighbour are found. Fragments below *min_size* voxels are dropped.
169    """
170    distance_3d = ndimage.distance_transform_edt(cell_mask, sampling=sampling)
171    markers = np.zeros(cell_mask.shape, dtype="int32")
172    next_id = 1
173    for z in range(cell_mask.shape[0]):
174        section = cell_mask[z]
175        if not section.any():
176            continue
177        distance = ndimage.distance_transform_edt(section)
178        peaks = peak_local_max(
179            distance, min_distance=min_distance, labels=connected_components(section), exclude_border=False
180        )
181        markers[z][tuple(peaks.T)] = np.arange(next_id, next_id + len(peaks))
182        next_id += len(peaks)
183    segments = watershed(-distance_3d, markers, mask=cell_mask)
184    instances = _merge_without_neck(segments, distance_3d, max_saddle_ratio, min_peak=min_distance // 4)
185    ids, counts = np.unique(instances, return_counts=True)
186    instances[np.isin(instances, ids[(counts < min_size) & (ids > 0)])] = 0
187    return np.unique(instances, return_inverse=True)[1].reshape(cell_mask.shape).astype("uint32")
188
189
190def _add_cell_instances(data_path):
191    """Write the derived cell instances into the h5 file once."""
192    import h5py
193
194    with h5py.File(data_path, "r") as f:
195        if CELL_INSTANCE_KEY in f:
196            return
197        cell_mask = f["labels/cell"][:] > 0
198    instances = _derive_cell_instances(cell_mask)
199    with h5py.File(data_path, "a") as f:
200        f.create_dataset(CELL_INSTANCE_KEY, data=instances, compression="gzip")
201
202
203def get_densecell_paths(
204    path: Union[os.PathLike, str], split: Literal["train", "val", "test"], download: bool = False
205) -> str:
206    """Get paths to the DenseCell data.
207
208    Args:
209        path: Filepath to a folder where the downloaded data will be saved.
210        split: The data split. Either 'train', 'val', or 'test'.
211        download: Whether to download the data if it is not present.
212
213    Returns:
214        The filepath for the stored data.
215    """
216    get_densecell_data(path, split, download)
217    data_path = os.path.join(path, f"densecell_{split}.h5")
218    return data_path
219
220
221def get_densecell_dataset(
222    path: Union[os.PathLike, str],
223    split: Literal["train", "val", "test"],
224    patch_shape: Tuple[int, int, int],
225    label_choice: Optional[str] = None,
226    download: bool = False,
227    **kwargs
228) -> Dataset:
229    """Get dataset for segmentation of organelles in SBF-SEM platelet images.
230
231    Args:
232        path: Filepath to a folder where the downloaded data will be saved.
233        split: The data split. Either 'train', 'val', or 'test'.
234        patch_shape: The patch shape to use for training.
235        label_choice: The organelle to segment. Available choices are:
236            'cell', 'mitochondrion', 'alpha_granule', 'canalicular_vessel', 'dense_granule', 'dense_core', or
237            'cell_instances', which splits the binary cell mask into instances with a 3D distance-transform
238            watershed and caches the result in the h5 file.
239            If None, uses 'original' which contains all semantic labels (0-6).
240        download: Whether to download the data if it is not present.
241        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
242
243    Returns:
244        The segmentation dataset.
245    """
246    assert split in ("train", "val", "test")
247
248    data_path = get_densecell_paths(path, split, download)
249
250    if label_choice is None:
251        label_key = "labels/original"
252    elif label_choice == "cell_instances":
253        _add_cell_instances(data_path)
254        label_key = CELL_INSTANCE_KEY
255    else:
256        valid_choices = list(ORGANELLES.values())
257        assert label_choice in valid_choices, f"'{label_choice}' is not valid. Choose from {valid_choices}."
258        label_key = f"labels/{label_choice}"
259
260    return torch_em.default_segmentation_dataset(
261        raw_paths=data_path,
262        raw_key="raw",
263        label_paths=data_path,
264        label_key=label_key,
265        patch_shape=patch_shape,
266        **kwargs
267    )
268
269
270def get_densecell_loader(
271    path: Union[os.PathLike, str],
272    split: Literal["train", "val", "test"],
273    patch_shape: Tuple[int, int, int],
274    batch_size: int,
275    label_choice: Optional[str] = None,
276    download: bool = False,
277    **kwargs
278) -> DataLoader:
279    """Get dataloader for segmentation of organelles in SBF-SEM platelet images.
280
281    Args:
282        path: Filepath to a folder where the downloaded data will be saved.
283        split: The data split. Either 'train', 'val', or 'test'.
284        patch_shape: The patch shape to use for training.
285        batch_size: The batch size for training.
286        label_choice: The organelle to segment. Available choices are:
287            'cell', 'mitochondrion', 'alpha_granule', 'canalicular_vessel', 'dense_granule', 'dense_core', or
288            'cell_instances', which splits the binary cell mask into instances with a 3D distance-transform
289            watershed and caches the result in the h5 file.
290            If None, uses 'original' which contains all semantic labels (0-6).
291        download: Whether to download the data if it is not present.
292        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
293
294    Returns:
295        The PyTorch DataLoader.
296    """
297    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
298    dataset = get_densecell_dataset(path, split, patch_shape, label_choice=label_choice, download=download, **ds_kwargs)
299    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)