torch_em.data.datasets.light_microscopy.neisseria_meningitidis

  1"""This dataset contains spinning-disk confocal time-lapse images of growing *Neisseria meningitidis*
  2bacterial colonies, with per-cell instance segmentation and lineage tracking (including division events).
  3The segmentation and tracks are reconstructed here from a TrackMate ilastik-detector tracking session
  4shipped with the raw data, by rasterizing the per-spot contours and assigning a new instance id to each
  5daughter cell after a division.
  6
  7NOTE: The instance contours come from an ilastik pixel classifier (used as the TrackMate detector), not
  8from manual annotation. They are comparatively rough and not densely accurate at the pixel level.
  9
 10The dataset is hosted on Zenodo at https://doi.org/10.5281/zenodo.5419619.
 11
 12Please cite it if you use this dataset for your research.
 13"""
 14
 15import os
 16import xml.etree.ElementTree as ET
 17from collections import defaultdict, deque
 18from typing import Dict, Tuple, Union
 19
 20import numpy as np
 21import tifffile
 22from skimage.draw import polygon
 23
 24from torch.utils.data import Dataset, DataLoader
 25
 26import torch_em
 27
 28from .. import util
 29
 30
 31URLS = {
 32    "raw": "https://zenodo.org/records/5419619/files/NeisseriaMeningitidisGrowth.tif",
 33    "tracks": "https://zenodo.org/records/5419619/files/NeisseriaMeningitidisGrowth.xml",
 34}
 35CHECKSUMS = {
 36    "raw": "491212b547654b0637001ce61e01cf289b63704ced7173b9dc44f2189fdc2ba9",
 37    "tracks": "a860f67520f7f5435be2d2d0a1759906e6f10c7dddffbad8ca7ce3ef24dd1a3c",
 38}
 39
 40
 41def get_neisseria_meningitidis_data(path: Union[os.PathLike, str], download: bool = False) -> Tuple[str, str]:
 42    """Download the Neisseria meningitidis bacterial growth dataset.
 43
 44    Args:
 45        path: Filepath to a folder where the downloaded data will be saved.
 46        download: Whether to download the data if it is not present.
 47
 48    Returns:
 49        Filepath to the raw image stack.
 50        Filepath to the TrackMate tracking / lineage annotations.
 51    """
 52    os.makedirs(path, exist_ok=True)
 53
 54    raw_path = os.path.join(path, "NeisseriaMeningitidisGrowth.tif")
 55    util.download_source(path=raw_path, url=URLS["raw"], download=download, checksum=CHECKSUMS["raw"])
 56
 57    tracks_path = os.path.join(path, "NeisseriaMeningitidisGrowth.xml")
 58    util.download_source(path=tracks_path, url=URLS["tracks"], download=download, checksum=CHECKSUMS["tracks"])
 59
 60    return raw_path, tracks_path
 61
 62
 63def _build_segment_ids(model) -> Dict[str, int]:
 64    # Assign a new instance id at each lineage root and to each daughter after a division, so that a
 65    # continuous single-cell trajectory keeps one id until it splits (CTC-style tracking convention).
 66    all_tracks = model.find("AllTracks")
 67    filtered_ids = {t.attrib["TRACK_ID"] for t in model.find("FilteredTracks").findall("TrackID")}
 68
 69    spot_frame = {}
 70    for frame_elem in model.find("AllSpots").findall("SpotsInFrame"):
 71        frame = int(frame_elem.attrib["frame"])
 72        for spot in frame_elem.findall("Spot"):
 73            spot_frame[spot.attrib["ID"]] = frame
 74
 75    out_edges, in_edges, kept_spots = defaultdict(list), defaultdict(list), set()
 76    for track in all_tracks.findall("Track"):
 77        if track.attrib["TRACK_ID"] not in filtered_ids:
 78            continue
 79        for edge in track.findall("Edge"):
 80            source, target = edge.attrib["SPOT_SOURCE_ID"], edge.attrib["SPOT_TARGET_ID"]
 81            if spot_frame[source] > spot_frame[target]:
 82                source, target = target, source
 83            out_edges[source].append(target)
 84            in_edges[target].append(source)
 85            kept_spots.update((source, target))
 86
 87    segment_id, next_id = {}, 1
 88    roots = sorted((s for s in kept_spots if not in_edges[s]), key=lambda s: spot_frame[s])
 89    queue = deque()
 90    for root in roots:
 91        segment_id[root] = next_id
 92        next_id += 1
 93        queue.append(root)
 94
 95    while queue:
 96        spot = queue.popleft()
 97        children = out_edges.get(spot, [])
 98        if len(children) == 1:
 99            segment_id[children[0]] = segment_id[spot]
100            queue.append(children[0])
101        else:
102            for child in children:
103                segment_id[child] = next_id
104                next_id += 1
105                queue.append(child)
106
107    return segment_id
108
109
110def _rasterize_labels(tracks_path: str, shape: Tuple[int, int, int]) -> np.ndarray:
111    root = ET.parse(tracks_path).getroot()
112    model = root.find("Model")
113    pixel_size = float(root.find("Settings").find("ImageData").attrib["pixelwidth"])
114    segment_id = _build_segment_ids(model)
115
116    _, height, width = shape
117    labels = np.zeros(shape, dtype=np.uint16)
118    for frame_elem in model.find("AllSpots").findall("SpotsInFrame"):
119        frame = int(frame_elem.attrib["frame"])
120        for spot in frame_elem.findall("Spot"):
121            label = segment_id.get(spot.attrib["ID"])
122            if label is None:  # Spots outside the filtered tracks are discarded, not segmented.
123                continue
124            center = np.array([float(spot.attrib["POSITION_X"]), float(spot.attrib["POSITION_Y"])])
125            points = np.array((spot.text or "").split(), dtype=float).reshape(-1, 2)
126            coords = (points + center) / pixel_size
127            rows, cols = polygon(coords[:, 1], coords[:, 0], shape=(height, width))
128            labels[frame, rows, cols] = label
129
130    return labels
131
132
133def get_neisseria_meningitidis_paths(path: Union[os.PathLike, str], download: bool = False) -> Tuple[str, str]:
134    """Get paths for the Neisseria meningitidis bacterial growth dataset.
135
136    Args:
137        path: Filepath to a folder where the downloaded data will be saved.
138        download: Whether to download the data if it is not present.
139
140    Returns:
141        Filepath to the raw image stack.
142        Filepath to the instance segmentation and tracking labels.
143    """
144    raw_path, tracks_path = get_neisseria_meningitidis_data(path, download)
145
146    label_path = os.path.join(path, "NeisseriaMeningitidisGrowth_labels.tif")
147    if not os.path.exists(label_path):
148        raw_shape = tifffile.TiffFile(raw_path).series[0].shape
149        tifffile.imwrite(label_path, _rasterize_labels(tracks_path, raw_shape))
150
151    return raw_path, label_path
152
153
154def get_neisseria_meningitidis_dataset(
155    path: Union[os.PathLike, str], patch_shape: Tuple[int, int, int], download: bool = False, **kwargs
156) -> Dataset:
157    """Get the Neisseria meningitidis dataset for bacterial cell segmentation and tracking.
158
159    Args:
160        path: Filepath to a folder where the downloaded data will be saved.
161        patch_shape: The patch shape to use for training.
162        download: Whether to download the data if it is not present.
163        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
164
165    Returns:
166        The segmentation dataset.
167    """
168    raw_path, label_path = get_neisseria_meningitidis_paths(path, download)
169
170    kwargs = util.update_kwargs(kwargs, "ndim", 2)
171
172    return torch_em.default_segmentation_dataset(
173        raw_paths=raw_path,
174        raw_key=None,
175        label_paths=label_path,
176        label_key=None,
177        patch_shape=patch_shape,
178        is_seg_dataset=True,
179        **kwargs
180    )
181
182
183def get_neisseria_meningitidis_loader(
184    path: Union[os.PathLike, str],
185    batch_size: int,
186    patch_shape: Tuple[int, int, int],
187    download: bool = False,
188    **kwargs
189) -> DataLoader:
190    """Get the Neisseria meningitidis dataloader for bacterial cell segmentation and tracking.
191
192    Args:
193        path: Filepath to a folder where the downloaded data will be saved.
194        batch_size: The batch size for training.
195        patch_shape: The patch shape to use for training.
196        download: Whether to download the data if it is not present.
197        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
198
199    Returns:
200        The DataLoader.
201    """
202    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
203    dataset = get_neisseria_meningitidis_dataset(path, patch_shape, download, **ds_kwargs)
204    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)