torch_em.data.datasets.light_microscopy.cellular

  1"""The CELLULAR dataset contains annotations for cell segmentation and autophagy classification
  2in fluorescence microscopy images of Drosophila cells.
  3
  4The cells express the autophagy reporter mRFP-EGFP-Atg8a. The EGFP signal fades once an
  5autophagosome fuses with a lysosome, while the mRFP signal stays, so the ratio of the two channels
  6shows whether autophagy is active. An expert marked every cell with one of three classes:
  7- 'fed' for basal autophagy,
  8- 'unfed' for activated autophagy,
  9- 'Unidentified' for a cell that the expert could not assign.
 10
 11The dataset holds 14005 annotated cells over 53 fields of view. Every field comes as one h5 file
 12with the three raw channels, an instance label image and a semantic label image.
 13
 14NOTE: The raw data has three channels, which the microscope metadata names FITC, Texas Red and
 15TL 25. This loader stores them as 'raw/fitc' for the EGFP signal, 'raw/texas_red' for the mRFP
 16signal and 'raw/brightfield' for the transmitted light.
 17
 18NOTE: The archive stores every cell as its own image of the full field, so a field of 264 cells
 19holds 264 images of 2048 x 2048 pixels. The preprocessing therefore reads about 170 GB of image
 20data to build 53 label images, and it runs over several cores.
 21
 22NOTE: The class of a cell comes from the folder that holds its mask. The masks are colored, but the
 23color varies inside one class, so it does not encode the class.
 24
 25NOTE: Only 53 of the 6240 fields carry masks. The loader reads the members of those fields out of
 26the archives, so it transfers about 1.4 GB instead of the full 110 GB.
 27
 28NOTE: Almost half of the cells are 'Unidentified'. Treat that class as unlabelled rather than as a
 29third biological state.
 30
 31The dataset is located at https://doi.org/10.5281/zenodo.8315423 under the CC BY 4.0 license.
 32This dataset is from the publication https://doi.org/10.1038/s41597-023-02687-x.
 33Please cite it if you use this dataset in your research.
 34"""
 35
 36import os
 37import zipfile
 38from glob import glob
 39from natsort import natsorted
 40from concurrent.futures import ProcessPoolExecutor
 41from typing import List, Literal, Optional, Sequence, Tuple, Union
 42
 43import numpy as np
 44import imageio.v3 as imageio
 45
 46from torch.utils.data import DataLoader, Dataset
 47
 48import torch_em
 49
 50from .. import util
 51
 52
 53URLS = {
 54    "images": "https://zenodo.org/records/8315423/files/images.zip?download=1",
 55    "masks": "https://zenodo.org/records/8315423/files/masks.zip?download=1",
 56}
 57
 58CHECKSUMS = {
 59    "images": None,  # The archive holds 105 GB, and the loader reads 159 members out of it.
 60    "masks": "8e640ba69e627e223a3703a883b5dda426bbbd069dddd61a5122f447bf880987",
 61}
 62
 63# The raw file of a channel ends with this token, and the microscope metadata names the channel.
 64CHANNELS = {"w1": "fitc", "w2": "texas_red", "w3": "brightfield"}
 65
 66CLASS_IDS = {"fed": 1, "unfed": 2, "Unidentified": 3}
 67
 68RAW_KEYS = tuple(f"raw/{name}" for name in CHANNELS.values())
 69
 70
 71def _list_fields(mask_dir: str) -> List[str]:
 72    """List the fields of view that carry masks."""
 73    return natsorted(os.path.basename(p) for p in glob(os.path.join(mask_dir, "*")) if os.path.isdir(p))
 74
 75
 76def _fetch_members(url: str, members: Sequence[str], destination: str) -> None:
 77    """Read the given members out of a remote zip archive, without downloading all of it."""
 78    import fsspec
 79
 80    with zipfile.ZipFile(fsspec.open(url, "rb").open()) as archive:
 81        available = set(archive.namelist())
 82        missing = [name for name in members if name not in available]
 83        if missing:
 84            raise RuntimeError(f"The archive {url} misses {len(missing)} members, e.g. '{missing[0]}'.")
 85        archive.extractall(destination, members=list(members))
 86
 87
 88def _build_field(arguments) -> Tuple[str, int]:
 89    """Write the h5 file of one field, and report how many cells it holds."""
 90    import h5py
 91    import tifffile
 92
 93    data_dir, field = arguments
 94    output_path = os.path.join(data_dir, "preprocessed", f"{field}.h5")
 95    if os.path.exists(output_path):
 96        return field, 0
 97
 98    raw = {
 99        name: tifffile.imread(os.path.join(data_dir, "images", f"{field}_{channel}.TIF"))
100        for channel, name in CHANNELS.items()
101    }
102    shape = raw["fitc"].shape
103
104    instances = np.zeros(shape, dtype="uint16")
105    semantic = np.zeros(shape, dtype="uint8")
106
107    instance_id = 0
108    for class_name, class_id in CLASS_IDS.items():
109        mask_paths = natsorted(glob(os.path.join(data_dir, "masks", field, class_name, "*.png")))
110        for mask_path in mask_paths:
111            mask = imageio.imread(mask_path)
112            mask = mask.max(axis=-1) > 0 if mask.ndim == 3 else mask > 0
113            # One mask of the archive holds no object at all.
114            if not mask.any():
115                continue
116            instance_id += 1
117            instances[mask] = instance_id
118            semantic[mask] = class_id
119
120    # Write to a temporary name, so that an interrupted run leaves no half written file behind.
121    temporary_path = f"{output_path}.tmp"
122    with h5py.File(temporary_path, "w") as f:
123        for name, array in raw.items():
124            f.create_dataset(f"raw/{name}", data=array, compression="gzip")
125        f.create_dataset("labels/instances", data=instances, compression="gzip")
126        f.create_dataset("labels/semantic", data=semantic, compression="gzip")
127    os.replace(temporary_path, output_path)
128
129    return field, instance_id
130
131
132def _preprocess(data_dir: str, n_workers: Optional[int] = None) -> str:
133    """Pack every field into one h5 file. The fields are independent, so they run in parallel."""
134    from tqdm import tqdm
135
136    output_dir = os.path.join(data_dir, "preprocessed")
137    os.makedirs(output_dir, exist_ok=True)
138
139    fields = _list_fields(os.path.join(data_dir, "masks"))
140    todo = [f for f in fields if not os.path.exists(os.path.join(output_dir, f"{f}.h5"))]
141    if not todo:
142        return output_dir
143
144    if n_workers is None:
145        n_workers = max(1, min(12, (os.cpu_count() or 1) - 2))
146
147    arguments = [(data_dir, field) for field in todo]
148    with ProcessPoolExecutor(n_workers) as executor:
149        list(tqdm(executor.map(_build_field, arguments), total=len(todo), desc="Preprocess the fields"))
150
151    return output_dir
152
153
154def get_cellular_data(
155    path: Union[os.PathLike, str], download: bool = False, n_workers: Optional[int] = None,
156) -> str:
157    """Download the CELLULAR dataset.
158
159    The loader reads the 53 annotated fields out of the archives, so it transfers about 1.4 GB
160    instead of the 110 GB that the repository holds.
161
162    Args:
163        path: Filepath to a folder where the downloaded data will be saved.
164        download: Whether to download the data if it is not present.
165        n_workers: The number of processes for the preprocessing. Defaults to the core count.
166
167    Returns:
168        The filepath to the folder with the h5 files.
169    """
170    data_dir = str(path)
171    output_dir = os.path.join(data_dir, "preprocessed")
172    if os.path.exists(output_dir) and glob(os.path.join(output_dir, "*.h5")):
173        return output_dir
174
175    if not download:
176        raise RuntimeError(f"Cannot find the data at {data_dir}, but download was set to False.")
177
178    os.makedirs(data_dir, exist_ok=True)
179
180    mask_dir = os.path.join(data_dir, "masks")
181    if not os.path.exists(mask_dir):
182        zip_path = os.path.join(data_dir, "masks.zip")
183        util.download_source(zip_path, URLS["masks"], download, CHECKSUMS["masks"])
184        util.unzip(zip_path=zip_path, dst=data_dir)
185
186    fields = _list_fields(mask_dir)
187    image_dir = os.path.join(data_dir, "images")
188    if not os.path.exists(image_dir):
189        members = [f"images/{field}_{channel}.TIF" for field in fields for channel in CHANNELS]
190        _fetch_members(URLS["images"], members, data_dir)
191
192    return _preprocess(data_dir, n_workers)
193
194
195def get_cellular_paths(
196    path: Union[os.PathLike, str],
197    fields: Optional[Sequence[str]] = None,
198    download: bool = False,
199) -> List[str]:
200    """Get paths to the CELLULAR data.
201
202    Args:
203        path: Filepath to a folder where the downloaded data will be saved.
204        fields: The fields of view to use. Defaults to all of them.
205        download: Whether to download the data if it is not present.
206
207    Returns:
208        List of filepaths for the h5 data.
209    """
210    output_dir = get_cellular_data(path, download)
211    volume_paths = natsorted(glob(os.path.join(output_dir, "*.h5")))
212
213    if fields is not None:
214        wanted = set(fields)
215        volume_paths = [p for p in volume_paths if os.path.splitext(os.path.basename(p))[0] in wanted]
216
217    if not volume_paths:
218        raise RuntimeError(f"Could not find any CELLULAR data in {output_dir}.")
219
220    return volume_paths
221
222
223def get_cellular_dataset(
224    path: Union[os.PathLike, str],
225    patch_shape: Tuple[int, int],
226    fields: Optional[Sequence[str]] = None,
227    channel: Literal["fitc", "texas_red", "brightfield", "all"] = "all",
228    label_choice: Literal["instances", "semantic"] = "instances",
229    offsets: Optional[List[List[int]]] = None,
230    boundaries: bool = False,
231    binary: bool = False,
232    download: bool = False,
233    **kwargs,
234) -> Dataset:
235    """Get the CELLULAR dataset for cell segmentation.
236
237    Args:
238        path: Filepath to a folder where the downloaded data will be saved.
239        patch_shape: The 2D patch shape to use for training.
240        fields: The fields of view to use. Defaults to all of them.
241        channel: The raw channel. Either 'fitc' for the EGFP signal, 'texas_red' for the mRFP
242            signal, 'brightfield' for the transmitted light, or 'all' for the three of them.
243        label_choice: The target. Either 'instances' for the single cells, or 'semantic' for the
244            autophagy classes, where one is fed, two is unfed and three is unidentified.
245        offsets: Offset values for affinity computation used as target.
246        boundaries: Whether to compute boundaries as the target.
247        binary: Whether to use a binary segmentation target.
248        download: Whether to download the data if it is not present.
249        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
250
251    Returns:
252        The segmentation dataset.
253    """
254    if len(patch_shape) != 2:
255        raise ValueError(f"The CELLULAR patch shape must be two-dimensional, got {patch_shape}.")
256    if label_choice not in ("instances", "semantic"):
257        raise ValueError(f"'{label_choice}' is not a valid label choice. Choose 'instances' or 'semantic'.")
258
259    valid_channels = tuple(CHANNELS.values()) + ("all",)
260    if channel not in valid_channels:
261        raise ValueError(f"'{channel}' is not a valid channel. Choose from {list(valid_channels)}.")
262
263    volume_paths = get_cellular_paths(path, fields, download)
264    raw_key = list(RAW_KEYS) if channel == "all" else f"raw/{channel}"
265    label_key = f"labels/{label_choice}"
266
267    if label_choice == "instances":
268        kwargs, _ = util.add_instance_label_transform(
269            kwargs, add_binary_target=True, offsets=offsets, boundaries=boundaries, binary=binary,
270        )
271    kwargs = util.ensure_transforms(ndim=2, **kwargs)
272
273    return torch_em.default_segmentation_dataset(
274        raw_paths=volume_paths,
275        raw_key=raw_key,
276        label_paths=volume_paths,
277        label_key=label_key,
278        patch_shape=patch_shape,
279        ndim=2,
280        with_channels=channel == "all",
281        **kwargs,
282    )
283
284
285def get_cellular_loader(
286    path: Union[os.PathLike, str],
287    batch_size: int,
288    patch_shape: Tuple[int, int],
289    fields: Optional[Sequence[str]] = None,
290    channel: Literal["fitc", "texas_red", "brightfield", "all"] = "all",
291    label_choice: Literal["instances", "semantic"] = "instances",
292    offsets: Optional[List[List[int]]] = None,
293    boundaries: bool = False,
294    binary: bool = False,
295    download: bool = False,
296    **kwargs,
297) -> DataLoader:
298    """Get the CELLULAR dataloader for cell segmentation.
299
300    Args:
301        path: Filepath to a folder where the downloaded data will be saved.
302        batch_size: The batch size for training.
303        patch_shape: The 2D patch shape to use for training.
304        fields: The fields of view to use. Defaults to all of them.
305        channel: The raw channel. Either 'fitc', 'texas_red', 'brightfield' or 'all'.
306        label_choice: The target. Either 'instances' or 'semantic'.
307        offsets: Offset values for affinity computation used as target.
308        boundaries: Whether to compute boundaries as the target.
309        binary: Whether to use a binary segmentation target.
310        download: Whether to download the data if it is not present.
311        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or the PyTorch DataLoader.
312
313    Returns:
314        The DataLoader.
315    """
316    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
317    dataset = get_cellular_dataset(
318        path=path,
319        patch_shape=patch_shape,
320        fields=fields,
321        channel=channel,
322        label_choice=label_choice,
323        offsets=offsets,
324        boundaries=boundaries,
325        binary=binary,
326        download=download,
327        **ds_kwargs,
328    )
329    return torch_em.get_data_loader(dataset, batch_size=batch_size, **loader_kwargs)