torch_em.data.datasets.light_microscopy.ecoli_microcolony_lineage

  1"""This dataset contains phase-contrast time-lapse images of growing *E. coli* microcolonies with
  2per-frame single-cell instance segmentation and full lineage tracking (Schnitzcells format), covering
  3eight genetic pathways (toxin production, SOS-stress response and metabolism).
  4
  5The dataset is hosted on Zenodo at https://doi.org/10.5281/zenodo.268921.
  6The dataset is from the publication https://doi.org/10.1016/j.cels.2018.03.009.
  7
  8Please cite it if you use this dataset for your research.
  9"""
 10
 11import os
 12import re
 13from glob import glob
 14from typing import List, Optional, Tuple, Union
 15
 16import h5py
 17import numpy as np
 18import tifffile
 19import scipy.io as sio
 20from skimage.feature import match_template
 21
 22from torch.utils.data import Dataset, DataLoader
 23
 24import torch_em
 25
 26from .. import util
 27
 28
 29GENES = ["cib", "crosstalk", "metA", "pheA", "recA", "rpsM", "SOSInteraction", "trpL"]
 30
 31URLS = {gene: f"https://zenodo.org/records/268921/files/{gene}.zip" for gene in GENES}
 32
 33CHECKSUMS = {
 34    "cib": "1e6408984f531347861dce04d9b4a4af1f51a6c85d2cf7f38149158632928555",
 35    "crosstalk": "dfd41c10be2f4271efaa79c1bea7c650b56769936186ec491565b5140e9a3d95",
 36    "metA": "c12142c15eb88c13f9e9456de02e2bb5e7c9aaadbca6651315b19504cbf11397",
 37    "pheA": "e92ea5af5e4900f7c2498efed1478d7948a1c195579ec211dd633e08c8b06cea",
 38    "recA": "145bd270d574b6210e36a9a17ab6be1fb7214498e40a1118018b3c97f91dfe28",
 39    "rpsM": "e0c3c49f9f4d47772f7a50a17f722bdf9031a68999e9b0fab77247e58fbb845b",
 40    "SOSInteraction": "aafc6c5af6ac5ad5659d6e1ddc4827321d7497080333f1e86ccfbd8eb76bf0fc",
 41    "trpL": "022b0ca2afe94d1890251bcb7c70ee500806f8c247ed2ba07ede35ed60b017fb",
 42}
 43
 44
 45def _load_lineage_lookup(lin_path):
 46    lookup = {}
 47    try:
 48        lineage = sio.loadmat(lin_path, simplify_cells=True)["schnitzcells"]
 49        for track_id, schnitz in enumerate(lineage, start=1):
 50            frames = np.atleast_1d(schnitz["frames"])
 51            cellnos = np.atleast_1d(schnitz["cellno"])
 52            for frame, cellno in zip(frames, cellnos):
 53                lookup[(int(frame), int(cellno))] = track_id
 54    except NotImplementedError:  # MATLAB v7.3 (HDF5-based) files need h5py instead of scipy.io.
 55        with h5py.File(lin_path, "r") as f:
 56            frames_ds, cellno_ds = f["schnitzcells"]["frames"], f["schnitzcells"]["cellno"]
 57            for track_id in range(1, frames_ds.shape[0] + 1):
 58                frames = np.array(f[frames_ds[track_id - 1, 0]]).ravel()
 59                cellnos = np.array(f[cellno_ds[track_id - 1, 0]]).ravel()
 60                for frame, cellno in zip(frames, cellnos):
 61                    lookup[(int(frame), int(cellno))] = track_id
 62    return lookup
 63
 64
 65def _find_crop_offset(image, template, rect, pad=20):
 66    y0, x0, y1, x1 = rect
 67    h, w = template.shape
 68    wy0, wy1 = max(0, y0 - 1 - pad), min(image.shape[0], y1 + pad)
 69    wx0, wx1 = max(0, x0 - 1 - pad), min(image.shape[1], x1 + pad)
 70    window = image[wy0:wy1, wx0:wx1]
 71    if window.shape[0] < h or window.shape[1] < w:
 72        return y0 - 1, x0 - 1
 73    result = match_template(window.astype(np.float32), template.astype(np.float32))
 74    dy, dx = np.unravel_index(np.argmax(result), result.shape)
 75    return wy0 + dy, wx0 + dx
 76
 77
 78def _reconstruct_label_frame(seg_path, raw_shape, lineage_lookup, frame_num, raw_image):
 79    seg = sio.loadmat(seg_path, simplify_cells=True)
 80    rect = seg.get("rect")
 81    if rect is None or len(rect) != 4:
 82        return None
 83
 84    y0, x0, y1, x1 = [int(v) for v in rect]
 85    local_labels = seg["Lc"]
 86    phsub = seg["phsub"]
 87
 88    yy0, xx0 = _find_crop_offset(raw_image, phsub, (y0, x0, y1, x1))
 89    h, w = local_labels.shape
 90
 91    canvas = np.zeros(raw_shape, dtype=np.uint16)
 92    placed = np.zeros_like(local_labels, dtype=np.uint16)
 93    for local_id in np.unique(local_labels):
 94        if local_id == 0:
 95            continue
 96        global_id = lineage_lookup.get((frame_num, int(local_id)))
 97        if global_id is not None:
 98            placed[local_labels == local_id] = global_id
 99    canvas[yy0:yy0 + h, xx0:xx0 + w] = placed
100    return canvas
101
102
103def _prepare_colony_labels(colony_dir, cache_root):
104    colony_name = os.path.basename(colony_dir.rstrip("/"))
105    lin_path = os.path.join(colony_dir, "data", f"{colony_name}_lin.mat")
106    if not os.path.exists(lin_path):
107        return [], []
108
109    label_dir = os.path.join(cache_root, colony_name)
110    raw_paths_all = sorted(glob(os.path.join(colony_dir, "images", f"{colony_name}-p-*.tif")))
111    seg_paths = sorted(glob(os.path.join(colony_dir, "segmentation", f"{colony_name}seg*.mat")))
112    if not raw_paths_all or not seg_paths:
113        return [], []
114
115    seg_by_frame = {}
116    for seg_path in seg_paths:
117        fname = os.path.basename(seg_path)
118        match = re.search(r"seg(\d+)\.mat$", fname)
119        if match is None:
120            continue
121        seg_by_frame[int(match.group(1))] = seg_path
122
123    os.makedirs(label_dir, exist_ok=True)
124    lineage_lookup = _load_lineage_lookup(lin_path)
125
126    raw_paths, label_paths = [], []
127    for raw_path in raw_paths_all:
128        fname = os.path.basename(raw_path)
129        match = re.search(r"-p-(\d+)\.tif$", fname)
130        if match is None:
131            continue
132        frame_num = int(match.group(1))
133        seg_path = seg_by_frame.get(frame_num)
134        if seg_path is None:
135            continue
136
137        label_path = os.path.join(label_dir, fname)
138        if not os.path.exists(label_path):
139            raw_image = tifffile.imread(raw_path)
140            label = _reconstruct_label_frame(seg_path, raw_image.shape, lineage_lookup, frame_num, raw_image)
141            if label is None:
142                continue
143            tifffile.imwrite(label_path, label)
144
145        raw_paths.append(raw_path)
146        label_paths.append(label_path)
147
148    return raw_paths, label_paths
149
150
151def get_ecoli_microcolony_lineage_data(
152    path: Union[os.PathLike, str], genes: Optional[List[str]] = None, download: bool = False,
153) -> List[str]:
154    f"""Download the E. coli microcolony lineage dataset.
155
156    Args:
157        path: Filepath to a folder where the downloaded data will be saved.
158        genes: The genetic pathways to download. The available pathways are: {', '.join(GENES)}.
159            By default downloads all of them.
160        download: Whether to download the data if it is not present.
161
162    Returns:
163        List of filepaths to the folders where each pathway's data is stored.
164    """
165    genes = GENES if genes is None else genes
166    for gene in genes:
167        if gene not in GENES:
168            raise ValueError(f"'{gene}' is not a valid pathway, choose one of {GENES}.")
169
170    os.makedirs(path, exist_ok=True)
171
172    gene_dirs = []
173    for gene in genes:
174        gene_dir = os.path.join(path, gene)
175        if not os.path.exists(gene_dir):
176            zip_path = os.path.join(path, f"{gene}.zip")
177            util.download_source(path=zip_path, url=URLS[gene], download=download, checksum=CHECKSUMS[gene])
178            util.unzip(zip_path=zip_path, dst=path)
179        gene_dirs.append(gene_dir)
180
181    return gene_dirs
182
183
184def get_ecoli_microcolony_lineage_paths(
185    path: Union[os.PathLike, str], genes: Optional[List[str]] = None, download: bool = False,
186) -> Tuple[List[str], List[str]]:
187    """Get paths for the E. coli microcolony lineage dataset.
188
189    Args:
190        path: Filepath to a folder where the downloaded data will be saved.
191        genes: The genetic pathways to use. By default uses all of them.
192        download: Whether to download the data if it is not present.
193
194    Returns:
195        List of filepaths for the raw phase-contrast images.
196        List of filepaths for the reconstructed instance segmentation and lineage labels.
197    """
198    gene_dirs = get_ecoli_microcolony_lineage_data(path, genes, download)
199
200    raw_paths, label_paths = [], []
201    for gene_dir in gene_dirs:
202        cache_root = os.path.join(gene_dir, "labels_lineage")
203        colony_dirs = [d for d in glob(os.path.join(gene_dir, "*")) if os.path.isdir(d) and d != cache_root]
204        for colony_dir in colony_dirs:
205            this_raw_paths, this_label_paths = _prepare_colony_labels(colony_dir, cache_root)
206            raw_paths.extend(this_raw_paths)
207            label_paths.extend(this_label_paths)
208
209    assert raw_paths and len(raw_paths) == len(label_paths)
210    return raw_paths, label_paths
211
212
213def get_ecoli_microcolony_lineage_dataset(
214    path: Union[os.PathLike, str],
215    patch_shape: Tuple[int, int],
216    genes: Optional[List[str]] = None,
217    download: bool = False,
218    **kwargs
219) -> Dataset:
220    """Get the E. coli microcolony lineage dataset for cell segmentation.
221
222    Args:
223        path: Filepath to a folder where the downloaded data will be saved.
224        patch_shape: The patch shape to use for training.
225        genes: The genetic pathways to use. By default uses all of them.
226        download: Whether to download the data if it is not present.
227        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
228
229    Returns:
230        The segmentation dataset.
231    """
232    raw_paths, label_paths = get_ecoli_microcolony_lineage_paths(path, genes, download)
233
234    return torch_em.default_segmentation_dataset(
235        raw_paths=raw_paths,
236        raw_key=None,
237        label_paths=label_paths,
238        label_key=None,
239        patch_shape=patch_shape,
240        ndim=2,
241        is_seg_dataset=False,
242        **kwargs
243    )
244
245
246def get_ecoli_microcolony_lineage_loader(
247    path: Union[os.PathLike, str],
248    batch_size: int,
249    patch_shape: Tuple[int, int],
250    genes: Optional[List[str]] = None,
251    download: bool = False,
252    **kwargs
253) -> DataLoader:
254    """Get the E. coli microcolony lineage dataloader for cell segmentation.
255
256    Args:
257        path: Filepath to a folder where the downloaded data will be saved.
258        batch_size: The batch size for training.
259        patch_shape: The patch shape to use for training.
260        genes: The genetic pathways to use. By default uses all of them.
261        download: Whether to download the data if it is not present.
262        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
263
264    Returns:
265        The DataLoader.
266    """
267    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
268    dataset = get_ecoli_microcolony_lineage_dataset(path, patch_shape, genes, download, **ds_kwargs)
269    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)