torch_em.data.datasets.electron_microscopy.deepcontact

  1"""DeepContact dataset for organelle segmentation in 2D EM.
  2
  3The dataset contains 2D SEM and TEM images of cultured cells and tissue with
  4manual polygon annotations for three organelle classes:
  5- mito: mitochondria
  6- er: endoplasmic reticulum
  7- ld: lipid droplets
  8
  9Data is provided as LabelMe JSON annotations paired with TIFF images.
 10During preprocessing, polygon annotations are rasterized to binary label masks
 11and stored alongside the raw images in HDF5 files.
 12
 13Three image sources are available:
 14- cell: SEM images of U-2 OS cultured cells at 5 nm/px
 15- sem: SEM images of Sertoli tissue cells at 10 nm/px
 16- tem: TEM images at 4.68 nm/px
 17
 18This dataset is from the publication https://doi.org/10.1083/jcb.202106190.
 19Please cite it if you use this dataset in your research.
 20
 21The data is available at https://figshare.com/articles/dataset/DeepContact_Training_Data/19898404.
 22"""
 23
 24import os
 25from glob import glob
 26from typing import List, Literal, Optional, Tuple, Union
 27
 28import numpy as np
 29from tqdm import tqdm
 30
 31import torch_em
 32from torch.utils.data import Dataset, DataLoader
 33from .. import util
 34
 35
 36DEEPCONTACT_URLS = {
 37    "cell": ("https://ndownloader.figshare.com/files/35317564", "cell_data.zip"),
 38    "sem": ("https://ndownloader.figshare.com/files/35317573", "sem_data.zip"),
 39    "tem": ("https://ndownloader.figshare.com/files/35317576", "tem_data.zip"),
 40}
 41
 42DEEPCONTACT_CHECKSUMS = {
 43    "cell": None,
 44    "sem": None,
 45    "tem": None,
 46}
 47
 48DEEPCONTACT_LABEL_NAMES = {
 49    "mito": ["Mito", "mito", "Mitochondria", "mitochondria"],
 50    "er": ["ER", "er"],
 51    "ld": ["Lipid Droplets", "lipid droplets", "LipidDroplets", "LD", "ld"],
 52}
 53
 54
 55def _rasterize_labelme_json(json_path, label_choice):
 56    import json
 57    from skimage.draw import polygon as sk_polygon
 58
 59    with open(json_path) as f:
 60        data = json.load(f)
 61
 62    h = data.get("imageHeight") or data.get("image_height")
 63    w = data.get("imageWidth") or data.get("image_width")
 64    mask = np.zeros((h, w), dtype=np.uint8)
 65
 66    target_names = DEEPCONTACT_LABEL_NAMES[label_choice]
 67    for shape in data.get("shapes", []):
 68        if shape.get("label") not in target_names:
 69            continue
 70        pts = np.array(shape["points"])
 71        rr, cc = sk_polygon(pts[:, 1], pts[:, 0], shape=(h, w))
 72        mask[rr, cc] = 1
 73
 74    return mask
 75
 76
 77def _find_image_for_json(json_path):
 78    from imageio import imread
 79
 80    base = os.path.splitext(json_path)[0]
 81    for ext in [".tif", ".tiff", ".png", ".jpg", ".jpeg"]:
 82        img_path = base + ext
 83        if os.path.exists(img_path):
 84            return imread(img_path)
 85
 86    # Check imagePath field in JSON
 87    import json
 88    with open(json_path) as f:
 89        data = json.load(f)
 90    img_name = data.get("imagePath", "")
 91    img_path = os.path.join(os.path.dirname(json_path), img_name)
 92    if os.path.exists(img_path):
 93        return imread(img_path)
 94
 95    return None
 96
 97
 98def _preprocess_source(extract_dir, output_dir, source):
 99    from elf.io import open_file
100
101    os.makedirs(output_dir, exist_ok=True)
102    json_files = sorted(glob(os.path.join(extract_dir, "**", "*.json"), recursive=True))
103
104    for json_path in tqdm(json_files, desc=f"Processing {source}"):
105        name = os.path.splitext(os.path.relpath(json_path, extract_dir))[0].replace(os.sep, "_")
106        h5_path = os.path.join(output_dir, f"{name}.h5")
107        if os.path.exists(h5_path):
108            continue
109
110        raw = _find_image_for_json(json_path)
111        if raw is None:
112            continue
113
114        if raw.ndim == 3:
115            raw = raw[..., 0]
116
117        with open_file(h5_path, "a") as f:
118            f.create_dataset("raw", data=raw.astype(np.uint8), compression="gzip")
119            for label_choice in DEEPCONTACT_LABEL_NAMES:
120                mask = _rasterize_labelme_json(json_path, label_choice)
121                f.create_dataset(f"labels/{label_choice}", data=mask, compression="gzip")
122
123
124def get_deepcontact_data(
125    path: Union[os.PathLike, str],
126    sources: Optional[List[Literal["cell", "sem", "tem"]]] = None,
127    download: bool = False,
128) -> str:
129    """Download and preprocess the DeepContact dataset.
130
131    Args:
132        path: Filepath to a folder where the data will be saved.
133        sources: Which image sources to use. Defaults to all ("cell", "sem", "tem").
134        download: Whether to download the data if not present.
135
136    Returns:
137        Path to the folder containing preprocessed HDF5 files.
138    """
139    if sources is None:
140        sources = ["cell", "sem", "tem"]
141
142    processed_dir = os.path.join(str(path), "processed")
143    os.makedirs(str(path), exist_ok=True)
144
145    for source in sources:
146        source_dir = os.path.join(processed_dir, source)
147        if os.path.isdir(source_dir) and len(glob(os.path.join(source_dir, "*.h5"))) > 0:
148            continue
149
150        url, fname = DEEPCONTACT_URLS[source]
151        zip_path = os.path.join(str(path), fname)
152
153        if not os.path.exists(zip_path):
154            if not download:
155                raise RuntimeError(
156                    f"Data for source '{source}' not found at '{zip_path}'. "
157                    "Set download=True or download manually from "
158                    "https://figshare.com/articles/dataset/DeepContact_Training_Data/19898404."
159                )
160            util.download_source(zip_path, url, download, checksum=DEEPCONTACT_CHECKSUMS[source])
161
162        extract_dir = os.path.join(str(path), f"{source}_raw")
163        if not os.path.isdir(extract_dir):
164            util.unzip(zip_path, extract_dir, remove=False)
165
166        _preprocess_source(extract_dir, source_dir, source)
167
168    return processed_dir
169
170
171def get_deepcontact_paths(
172    path: Union[os.PathLike, str],
173    sources: Optional[List[Literal["cell", "sem", "tem"]]] = None,
174    download: bool = False,
175) -> List[str]:
176    """Get paths to DeepContact HDF5 files.
177
178    Args:
179        path: Filepath to a folder where the data will be saved.
180        sources: Which image sources to use. Defaults to all ("cell", "sem", "tem").
181        download: Whether to download the data if not present.
182
183    Returns:
184        List of paths to HDF5 files.
185    """
186    if sources is None:
187        sources = ["cell", "sem", "tem"]
188    processed_dir = get_deepcontact_data(path, sources, download)
189    paths = []
190    for source in sources:
191        paths.extend(sorted(glob(os.path.join(processed_dir, source, "*.h5"))))
192    return paths
193
194
195def get_deepcontact_dataset(
196    path: Union[os.PathLike, str],
197    patch_shape: Tuple[int, int],
198    label_choice: Literal["mito", "er", "ld"] = "mito",
199    sources: Optional[List[Literal["cell", "sem", "tem"]]] = None,
200    download: bool = False,
201    **kwargs,
202) -> Dataset:
203    """Get the DeepContact dataset for organelle segmentation in 2D EM.
204
205    Args:
206        path: Filepath to a folder where the data will be saved.
207        patch_shape: The patch shape (H, W) for training.
208        label_choice: Which organelle to segment. One of "mito", "er", or "ld".
209        sources: Which image sources to use. Defaults to all ("cell", "sem", "tem").
210        download: Whether to download the data if not present.
211        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
212
213    Returns:
214        The segmentation dataset.
215    """
216    assert len(patch_shape) == 2
217    data_paths = get_deepcontact_paths(path, sources, download)
218
219    return torch_em.default_segmentation_dataset(
220        raw_paths=data_paths,
221        raw_key="raw",
222        label_paths=data_paths,
223        label_key=f"labels/{label_choice}",
224        patch_shape=patch_shape,
225        **kwargs,
226    )
227
228
229def get_deepcontact_loader(
230    path: Union[os.PathLike, str],
231    batch_size: int,
232    patch_shape: Tuple[int, int],
233    label_choice: Literal["mito", "er", "ld"] = "mito",
234    sources: Optional[List[Literal["cell", "sem", "tem"]]] = None,
235    download: bool = False,
236    **kwargs,
237) -> DataLoader:
238    """Get the DataLoader for organelle segmentation in the DeepContact dataset.
239
240    Args:
241        path: Filepath to a folder where the data will be saved.
242        batch_size: The batch size for training.
243        patch_shape: The patch shape (H, W) for training.
244        label_choice: Which organelle to segment. One of "mito", "er", or "ld".
245        sources: Which image sources to use. Defaults to all ("cell", "sem", "tem").
246        download: Whether to download the data if not present.
247        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`
248            or for the PyTorch DataLoader.
249
250    Returns:
251        The DataLoader.
252    """
253    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
254    ds = get_deepcontact_dataset(
255        path=path,
256        patch_shape=patch_shape,
257        label_choice=label_choice,
258        sources=sources,
259        download=download,
260        **ds_kwargs,
261    )
262    return torch_em.get_data_loader(ds, batch_size=batch_size, **loader_kwargs)