torch_em.data.datasets.light_microscopy.alfi

  1"""The ALFI dataset contains annotations for cell segmentation in label-free
  2differential interference contrast timelapses of cultured human cells.
  3
  4ALFI stands for Annotations for Label-Free Images. The dataset holds 796 annotated frames of eight
  5timelapse sequences of the cell lines U2OS, HeLa and hTERT RPE-1. Every frame comes with a mask that
  6separates interphase cells from mitotic cells.
  7
  8The data suits nucleus segmentation in label-free bright-field microscopy, because the annotators
  9outlined the nucleus of every interphase cell. It also suits cell tracking, because the archive
 10holds a track id and a lineage parent for every annotated cell over time.
 11
 12NOTE: The two classes do not mark the same structure. The interphase label covers the nucleus, while
 13the mitotic label covers the whole rounded cell. The cytoplasm of an interphase cell is background.
 14So the target is a nucleus for an interphase cell, and a whole cell for a mitotic one.
 15
 16NOTE: This loader returns segmentation targets only. The tracking annotations sit next to the masks
 17in '<sequence>/<sequence>_DTLTruth.csv', which stores a frame index, a track id, a class, a bounding
 18box and the id of the parent cell. A second table, '<sequence>_PhenoTruth.csv', marks the phenotypes
 19early mitosis, late mitosis, cell death and multipolar division. Read these files directly if you
 20want to track cells or follow a lineage.
 21
 22NOTE: The archive holds 29 sequences, but only the eight MI sequences carry masks. The other
 23sequences provide bounding boxes, which this loader does not use. The loader reads the members of
 24the eight MI sequences out of the archive, so it transfers about 1.2 GB instead of the full 8.4 GB.
 25
 26NOTE: The publication defines no train, validation and test split. This loader splits by sequence,
 27so that a split never shares a sequence with another one. The frames are seven minutes apart and
 28look almost the same, so a split over single frames would leak.
 29
 30The dataset is located at https://doi.org/10.6084/m9.figshare.23798451 under the CC BY 4.0 license.
 31This dataset is from the publication https://doi.org/10.1038/s41597-023-02540-1.
 32Please cite it if you use this dataset in your research.
 33"""
 34
 35import os
 36import zipfile
 37from glob import glob
 38from natsort import natsorted
 39from typing import List, Literal, Optional, Sequence, Tuple, Union
 40
 41import numpy as np
 42import imageio.v3 as imageio
 43
 44from torch.utils.data import DataLoader, Dataset
 45
 46import torch_em
 47
 48from .. import util
 49
 50
 51URL = "https://ndownloader.figshare.com/files/41740227"
 52CHECKSUM = "fe3326323c10b1748302e962eae26150"
 53
 54# The folder name holds an ampersand, so quote it whenever it goes into a shell command.
 55ARCHIVE_ROOT = "Data&Annotations"
 56
 57SEQUENCES = ("MI01", "MI02", "MI03", "MI04", "MI05", "MI06", "MI07", "MI08")
 58
 59# The publication has no official split, so the sequences are grouped by cell line.
 60SPLITS = {
 61    "train": ("MI01", "MI02", "MI03", "MI04", "MI05"),  # U2OS
 62    "val": ("MI06",),  # HeLa
 63    "test": ("MI07", "MI08"),  # hTERT RPE-1
 64}
 65
 66# The masks store the background as 0, an interphase cell as 128 and a mitotic cell as 255.
 67SEMANTIC_IDS = {0: 0, 128: 1, 255: 2}
 68
 69# A few masks hold tiny blobs at the image border that the annotation tables do not list.
 70MIN_INSTANCE_SIZE = 200
 71
 72
 73def _extract_sequences(zip_path: str, path: str) -> None:
 74    """Extract the images and the masks of the annotated sequences."""
 75    with zipfile.ZipFile(zip_path) as archive:
 76        members = [
 77            name for name in archive.namelist()
 78            if name.startswith(f"{ARCHIVE_ROOT}/MI")
 79            and not name.startswith("__MACOSX")
 80            and not name.endswith(".DS_Store")
 81        ]
 82        if not members:
 83            raise RuntimeError(f"The archive {zip_path} does not hold the annotated MI sequences.")
 84        archive.extractall(path, members=members)
 85
 86
 87def _to_semantic(mask: np.ndarray) -> np.ndarray:
 88    """Map the mask values onto consecutive class ids."""
 89    semantic = np.zeros(mask.shape, dtype="uint8")
 90    for value, class_id in SEMANTIC_IDS.items():
 91        semantic[mask == value] = class_id
 92    return semantic
 93
 94
 95def _to_instances(mask: np.ndarray) -> np.ndarray:
 96    """Split each class of a mask into connected components and give every component one id."""
 97    from scipy.ndimage import label as connected_components
 98
 99    instances = np.zeros(mask.shape, dtype="uint16")
100    offset = 0
101    for value in (128, 255):
102        components, n_components = connected_components(mask == value)
103        for component_id in range(1, n_components + 1):
104            component = components == component_id
105            if component.sum() < MIN_INSTANCE_SIZE:
106                continue
107            offset += 1
108            instances[component] = offset
109    return instances
110
111
112def _create_labels(data_dir: str, sequence: str, label_choice: str) -> str:
113    """Convert the masks of one sequence into the requested target."""
114    from tqdm import tqdm
115
116    mask_dir = os.path.join(data_dir, sequence, "Masks")
117    label_dir = os.path.join(data_dir, sequence, f"{label_choice}_labels")
118    os.makedirs(label_dir, exist_ok=True)
119
120    mask_paths = natsorted(glob(os.path.join(mask_dir, "*.png")))
121    convert = _to_semantic if label_choice == "semantic" else _to_instances
122
123    for mask_path in tqdm(mask_paths, desc=f"Preprocess '{sequence}' for the {label_choice} target"):
124        output_path = os.path.join(label_dir, os.path.basename(mask_path).replace(".png", ".tif"))
125        if os.path.exists(output_path):
126            continue
127        imageio.imwrite(output_path, convert(imageio.imread(mask_path)), compression="zlib")
128
129    return label_dir
130
131
132def get_alfi_data(path: Union[os.PathLike, str], download: bool = False) -> str:
133    """Download the ALFI dataset.
134
135    The loader reads only the eight annotated sequences out of the archive, so it stores about
136    1.2 GB instead of the full 8.4 GB.
137
138    Args:
139        path: Filepath to a folder where the downloaded data will be saved.
140        download: Whether to download the data if it is not present.
141
142    Returns:
143        The filepath to the folder that holds the sequences.
144    """
145    data_dir = os.path.join(path, ARCHIVE_ROOT)
146    if os.path.exists(data_dir):
147        return data_dir
148
149    os.makedirs(path, exist_ok=True)
150    zip_path = os.path.join(path, "ALFIdatasetFinal.zip")
151    util.download_source(zip_path, URL, download, CHECKSUM)
152    _extract_sequences(zip_path, path)
153
154    return data_dir
155
156
157def get_alfi_paths(
158    path: Union[os.PathLike, str],
159    split: Optional[Literal["train", "val", "test"]] = "train",
160    sequences: Optional[Sequence[str]] = None,
161    label_choice: Literal["semantic", "instances"] = "instances",
162    download: bool = False,
163) -> Tuple[List[str], List[str]]:
164    """Get paths to the ALFI data.
165
166    Args:
167        path: Filepath to a folder where the downloaded data will be saved.
168        split: The data split, which groups the sequences by cell line. Either 'train' for U2OS,
169            'val' for HeLa or 'test' for hTERT RPE-1. Ignored when you pass `sequences`.
170        sequences: The sequences to use, for example ('MI01', 'MI06'). Overrides `split`.
171        label_choice: The target. Either 'instances' for the single cells, or 'semantic' for the
172            classes, where one is an interphase cell and two is a mitotic cell.
173        download: Whether to download the data if it is not present.
174
175    Returns:
176        List of filepaths for the image data.
177        List of filepaths for the label data.
178    """
179    if label_choice not in ("semantic", "instances"):
180        raise ValueError(f"'{label_choice}' is not a valid label choice. Choose 'semantic' or 'instances'.")
181
182    if sequences is None:
183        if split not in SPLITS:
184            raise ValueError(f"'{split}' is not a valid split. Choose from {list(SPLITS)}, or pass sequences.")
185        sequences = SPLITS[split]
186    else:
187        for sequence in sequences:
188            if sequence not in SEQUENCES:
189                raise ValueError(f"'{sequence}' is not an annotated sequence. Choose from {list(SEQUENCES)}.")
190
191    data_dir = get_alfi_data(path, download)
192
193    image_paths, label_paths = [], []
194    for sequence in sequences:
195        label_dir = _create_labels(data_dir, sequence, label_choice)
196        for label_path in natsorted(glob(os.path.join(label_dir, "*.tif"))):
197            # An image is named I_MI01_0001.png and its mask M_MI01_0001.png.
198            image_name = os.path.basename(label_path).replace(".tif", ".png").replace("M_", "I_", 1)
199            image_path = os.path.join(data_dir, sequence, "Images", image_name)
200            if not os.path.exists(image_path):
201                continue
202            image_paths.append(image_path)
203            label_paths.append(label_path)
204
205    if not image_paths:
206        raise RuntimeError(f"Could not find any ALFI data in {data_dir}.")
207
208    return image_paths, label_paths
209
210
211def get_alfi_dataset(
212    path: Union[os.PathLike, str],
213    patch_shape: Tuple[int, int],
214    split: Optional[Literal["train", "val", "test"]] = "train",
215    sequences: Optional[Sequence[str]] = None,
216    label_choice: Literal["semantic", "instances"] = "instances",
217    offsets: Optional[List[List[int]]] = None,
218    boundaries: bool = False,
219    binary: bool = False,
220    download: bool = False,
221    **kwargs,
222) -> Dataset:
223    """Get the ALFI dataset for cell segmentation.
224
225    Args:
226        path: Filepath to a folder where the downloaded data will be saved.
227        patch_shape: The 2D patch shape to use for training.
228        split: The data split, which groups the sequences by cell line. Ignored with `sequences`.
229        sequences: The sequences to use, for example ('MI01', 'MI06'). Overrides `split`.
230        label_choice: The target. Either 'instances' for the single cells, or 'semantic' for the
231            classes, where one is an interphase cell and two is a mitotic cell.
232        offsets: Offset values for affinity computation used as target.
233        boundaries: Whether to compute boundaries as the target.
234        binary: Whether to use a binary segmentation target.
235        download: Whether to download the data if it is not present.
236        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
237
238    Returns:
239        The segmentation dataset.
240    """
241    if len(patch_shape) != 2:
242        raise ValueError(f"The ALFI patch shape must be two-dimensional, got {patch_shape}.")
243
244    image_paths, label_paths = get_alfi_paths(path, split, sequences, label_choice, download)
245
246    if label_choice == "instances":
247        kwargs, _ = util.add_instance_label_transform(
248            kwargs, add_binary_target=True, offsets=offsets, boundaries=boundaries, binary=binary,
249        )
250    kwargs = util.ensure_transforms(ndim=2, **kwargs)
251
252    return torch_em.default_segmentation_dataset(
253        raw_paths=image_paths,
254        raw_key=None,
255        label_paths=label_paths,
256        label_key=None,
257        patch_shape=patch_shape,
258        is_seg_dataset=False,
259        ndim=2,
260        **kwargs,
261    )
262
263
264def get_alfi_loader(
265    path: Union[os.PathLike, str],
266    batch_size: int,
267    patch_shape: Tuple[int, int],
268    split: Optional[Literal["train", "val", "test"]] = "train",
269    sequences: Optional[Sequence[str]] = None,
270    label_choice: Literal["semantic", "instances"] = "instances",
271    offsets: Optional[List[List[int]]] = None,
272    boundaries: bool = False,
273    binary: bool = False,
274    download: bool = False,
275    **kwargs,
276) -> DataLoader:
277    """Get the ALFI dataloader for cell segmentation.
278
279    Args:
280        path: Filepath to a folder where the downloaded data will be saved.
281        batch_size: The batch size for training.
282        patch_shape: The 2D patch shape to use for training.
283        split: The data split, which groups the sequences by cell line. Ignored with `sequences`.
284        sequences: The sequences to use, for example ('MI01', 'MI06'). Overrides `split`.
285        label_choice: The target. Either 'instances' for the single cells, or 'semantic' for the
286            classes, where one is an interphase cell and two is a mitotic cell.
287        offsets: Offset values for affinity computation used as target.
288        boundaries: Whether to compute boundaries as the target.
289        binary: Whether to use a binary segmentation target.
290        download: Whether to download the data if it is not present.
291        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or the PyTorch DataLoader.
292
293    Returns:
294        The DataLoader.
295    """
296    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
297    dataset = get_alfi_dataset(
298        path=path,
299        patch_shape=patch_shape,
300        split=split,
301        sequences=sequences,
302        label_choice=label_choice,
303        offsets=offsets,
304        boundaries=boundaries,
305        binary=binary,
306        download=download,
307        **ds_kwargs,
308    )
309    return torch_em.get_data_loader(dataset, batch_size=batch_size, **loader_kwargs)