torch_em.data.datasets.light_microscopy.livecell

  1"""The LIVECell dataset contains phase-contrast microscopy images
  2and annotations for cell segmentations for 8 different cell lines.
  3
  4This dataset is described in the publication https://doi.org/10.1038/s41592-021-01249-6.
  5Please cite it if you use this dataset in your research.
  6"""
  7
  8import os
  9import json
 10import requests
 11from tqdm import tqdm
 12from shutil import copyfileobj
 13from typing import List, Optional, Sequence, Tuple, Union
 14
 15import numpy as np
 16import imageio.v3 as imageio
 17
 18import torch
 19from torch.utils.data import Dataset, DataLoader
 20
 21import torch_em
 22
 23from .. import util
 24from ... import ImageCollectionDataset
 25
 26try:
 27    from pycocotools.coco import COCO
 28except ImportError:
 29    COCO = None
 30
 31URLS = {
 32    "images": "http://livecell-dataset.s3.eu-central-1.amazonaws.com/LIVECell_dataset_2021/images.zip",
 33    "train": ("http://livecell-dataset.s3.eu-central-1.amazonaws.com/LIVECell_dataset_2021/annotations/"
 34              "LIVECell/livecell_coco_train.json"),
 35    "val": ("http://livecell-dataset.s3.eu-central-1.amazonaws.com/LIVECell_dataset_2021/annotations/"
 36            "LIVECell/livecell_coco_val.json"),
 37    "test": ("http://livecell-dataset.s3.eu-central-1.amazonaws.com/LIVECell_dataset_2021/annotations/"
 38             "LIVECell/livecell_coco_test.json")
 39}
 40# TODO
 41CHECKSUM = None
 42
 43CELL_TYPES = ['A172', 'BT474', 'BV2', 'Huh7', 'MCF7', 'SHSY5Y', 'SkBr3', 'SKOV3']
 44
 45
 46# TODO use download flag
 47def _download_annotation_file(path, split, download):
 48    annotation_file = os.path.join(path, f"{split}.json")
 49    if not os.path.exists(annotation_file):
 50        url = URLS[split]
 51        print("Downloading livecell annotation file from", url)
 52        with requests.get(url, stream=True) as r:
 53            with open(annotation_file, 'wb') as f:
 54                copyfileobj(r.raw, f)
 55    return annotation_file
 56
 57
 58def _annotations_to_instances(coco, image_metadata, category_ids):
 59    import bioimage_cpp as bic
 60
 61    # create and save the segmentation
 62    annotation_ids = coco.getAnnIds(imgIds=image_metadata["id"], catIds=category_ids)
 63    annotations = coco.loadAnns(annotation_ids)
 64    assert len(annotations) <= np.iinfo("uint16").max
 65    shape = (image_metadata["height"], image_metadata["width"])
 66    seg = np.zeros(shape, dtype="uint32")
 67
 68    # sort annotations by size, except for iscrowd which go first
 69    # we do this to minimize small noise from overlapping multi annotations
 70    # (see below)
 71    sizes = [ann["area"] if ann["iscrowd"] == 0 else 1 for ann in annotations]
 72    sorting = np.argsort(sizes)
 73    annotations = [annotations[i] for i in sorting]
 74
 75    for seg_id, annotation in enumerate(annotations, 1):
 76        mask = coco.annToMask(annotation).astype("bool")
 77        assert mask.shape == seg.shape
 78        seg[mask] = seg_id
 79
 80    # some images have multiple masks per object with slightly different foreground
 81    # this causes small noise objects we need to filter
 82    min_size = 50
 83    seg_ids, sizes = np.unique(seg, return_counts=True)
 84    seg[np.isin(seg, seg_ids[sizes < min_size])] = 0
 85
 86    seg, _, _ = bic.segmentation.relabel_sequential(seg)
 87
 88    return seg.astype("uint16")
 89
 90
 91def _create_segmentations_from_annotations(annotation_file, image_folder, seg_folder, cell_types, split):
 92    # Use a per-split and per-cell_types cache to avoid reloading the COCO JSON when data is already prepared.
 93    # The split must be part of the key: train and val share the same seg_folder, so a key without it
 94    # would return the paths of whichever split was processed first for both.
 95    cache_key = "all" if cell_types is None else "_".join(sorted(cell_types))
 96    cache_file = os.path.join(seg_folder, f"seg_paths_{split}_{cache_key}.json")
 97    if os.path.exists(cache_file):
 98        with open(cache_file) as f:
 99            cached = json.load(f)
100        image_paths = [os.path.join(seg_folder, fname) for fname in cached["image_paths"]]
101        seg_paths = [os.path.join(seg_folder, fname) for fname in cached["seg_paths"]]
102        return image_paths, seg_paths
103
104    if COCO is None:
105        raise ModuleNotFoundError(
106            "'pycocotools' is required for processing the LIVECell ground-truth. "
107            "Install it with 'conda install -c conda-forge pycocotools'."
108        )
109
110    coco = COCO(annotation_file)
111    category_ids = coco.getCatIds(catNms=["cell"])
112    image_ids = coco.getImgIds(catIds=category_ids)
113
114    image_paths, seg_paths = [], []
115    for image_id in tqdm(image_ids, desc="creating livecell segmentations from coco-style annotations"):
116        # get the path for the image data and make sure the corresponding image exists
117        image_metadata = coco.loadImgs(image_id)[0]
118        file_name = image_metadata["file_name"]
119
120        # if cell_type names are given we only select file names that match a cell_type
121        if cell_types is not None and (not any([cell_type in file_name for cell_type in cell_types])):
122            continue
123
124        sub_folder = file_name.split("_")[0]
125        image_path = os.path.join(image_folder, sub_folder, file_name)
126        # something changed in the image layout? we keep the old version around in case this changes back...
127        if not os.path.exists(image_path):
128            image_path = os.path.join(image_folder, file_name)
129        assert os.path.exists(image_path), image_path
130        image_paths.append(image_path)
131
132        # get the output path
133        out_folder = os.path.join(seg_folder, sub_folder)
134        os.makedirs(out_folder, exist_ok=True)
135        seg_path = os.path.join(out_folder, file_name)
136        seg_paths.append(seg_path)
137        if os.path.exists(seg_path):
138            continue
139
140        seg = _annotations_to_instances(coco, image_metadata, category_ids)
141        imageio.imwrite(seg_path, seg)
142
143    assert len(image_paths) == len(seg_paths)
144    assert len(image_paths) > 0, \
145        f"No matching image paths were found. Did you pass invalid cell type names ({cell_types})?"
146
147    cache_dir = os.path.dirname(cache_file)
148    image_paths_rel = [os.path.relpath(image_path, start=cache_dir) for image_path in image_paths]
149    seg_paths_rel = [os.path.relpath(seg_path, start=cache_dir) for seg_path in seg_paths]
150    with open(cache_file, "w") as f:
151        json.dump({"image_paths": image_paths_rel, "seg_paths": seg_paths_rel}, f)
152
153    return image_paths, seg_paths
154
155
156def _download_livecell_annotations(path, split, download, cell_types, label_path):
157    annotation_file = _download_annotation_file(path, split, download)
158    if split == "test":
159        split_name = "livecell_test_images"
160    else:
161        split_name = "livecell_train_val_images"
162
163    image_folder = os.path.join(path, "images", split_name)
164    seg_folder = os.path.join(path, "annotations", split_name) if label_path is None\
165        else os.path.join(label_path, "annotations", split_name)
166
167    assert os.path.exists(image_folder), image_folder
168
169    return _create_segmentations_from_annotations(annotation_file, image_folder, seg_folder, cell_types, split)
170
171
172def get_livecell_data(path: Union[os.PathLike], download: bool = False):
173    """Download the LIVECell dataset.
174
175    Args:
176        path: Filepath to a folder where the downloaded data will be saved.
177        download: Whether to download the data if it is not present.
178    """
179    os.makedirs(path, exist_ok=True)
180    image_path = os.path.join(path, "images")
181
182    if os.path.exists(image_path):
183        return
184
185    url = URLS["images"]
186    checksum = CHECKSUM
187    zip_path = os.path.join(path, "livecell.zip")
188    util.download_source(zip_path, url, download, checksum)
189    util.unzip(zip_path, path, True)
190
191
192def get_livecell_paths(
193    path: Union[os.PathLike, str],
194    split: str,
195    download: bool = False,
196    cell_types: Optional[Sequence[str]] = None,
197    label_path: Optional[Union[os.PathLike, str]] = None
198) -> Tuple[List[str], List[str]]:
199    """Get paths to the LIVECell data.
200
201    Args:
202        path: Filepath to a folder where the downloaded data will be saved.
203        split: The data split to use. Either 'train', 'val' or 'test'.
204        download: Whether to download the data if it is not present.
205        cell_types: The cell types for which to get the data paths.
206        label_path: Optional path for loading the label data.
207
208    Returns:
209        List of filepaths for the image data.
210        List of filepaths for the label data.
211    """
212    get_livecell_data(path, download)
213    image_paths, seg_paths = _download_livecell_annotations(path, split, download, cell_types, label_path)
214    return image_paths, seg_paths
215
216
217def get_livecell_dataset(
218    path: Union[os.PathLike, str],
219    split: str,
220    patch_shape: Tuple[int, int],
221    download: bool = False,
222    offsets: Optional[List[List[int]]] = None,
223    boundaries: bool = False,
224    binary: bool = False,
225    cell_types: Optional[Sequence[str]] = None,
226    label_path: Optional[Union[os.PathLike, str]] = None,
227    label_dtype=torch.int64,
228    **kwargs
229) -> Dataset:
230    """Get the LIVECell dataset for segmenting cells in phase-contrast microscopy.
231
232    Args:
233        path: Filepath to a folder where the downloaded data will be saved.
234        split: The data split to use. Either 'train', 'val' or 'test'.
235        patch_shape: The patch shape to use for training.
236        download: Whether to download the data if it is not present.
237        offsets: Offset values for affinity computation used as target.
238        boundaries: Whether to compute boundaries as the target.
239        binary: Whether to use a binary segmentation target.
240        cell_types: The cell types for which to get the data paths.
241        label_path: Optional path for loading the label data.
242        label_dtype: The datatype of the label data.
243        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
244
245    Returns:
246        The segmentation dataset.
247    """
248    assert split in ("train", "val", "test")
249    if cell_types is not None:
250        assert isinstance(cell_types, (list, tuple)), \
251            f"cell_types must be passed as a list or tuple instead of {cell_types}"
252
253    image_paths, seg_paths = get_livecell_paths(path, split, download, cell_types, label_path)
254
255    kwargs = util.ensure_transforms(ndim=2, **kwargs)
256    kwargs, label_dtype = util.add_instance_label_transform(
257        kwargs, add_binary_target=True, label_dtype=label_dtype, offsets=offsets, boundaries=boundaries, binary=binary
258    )
259
260    return ImageCollectionDataset(
261        raw_image_paths=image_paths,
262        label_image_paths=seg_paths,
263        patch_shape=patch_shape,
264        label_dtype=label_dtype,
265        **kwargs
266    )
267
268
269def get_livecell_loader(
270    path: Union[os.PathLike, str],
271    split: str,
272    patch_shape: Tuple[int, int],
273    batch_size: int,
274    download: bool = False,
275    offsets: Optional[List[List[int]]] = None,
276    boundaries: bool = False,
277    binary: bool = False,
278    cell_types: Optional[Sequence[str]] = None,
279    label_path: Optional[Union[os.PathLike, str]] = None,
280    label_dtype=torch.int64,
281    **kwargs
282) -> DataLoader:
283    """Get the LIVECell dataloader for segmenting cells in phase-contrast microscopy.
284
285    Args:
286        path: Filepath to a folder where the downloaded data will be saved.
287        split: The data split to use. Either 'train', 'val' or 'test'.
288        patch_shape: The patch shape to use for training.
289        batch_size: The batch size for training.
290        download: Whether to download the data if it is not present.
291        offsets: Offset values for affinity computation used as target.
292        boundaries: Whether to compute boundaries as the target.
293        binary: Whether to use a binary segmentation target.
294        cell_types: The cell types for which to get the data paths.
295        label_path: Optional path for loading the label data.
296        label_dtype: The datatype of the label data.
297        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
298
299    Returns:
300        The DataLoader.
301    """
302    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
303    dataset = get_livecell_dataset(
304        path, split, patch_shape, download=download, offsets=offsets, boundaries=boundaries, binary=binary,
305        cell_types=cell_types, label_path=label_path, label_dtype=label_dtype, **ds_kwargs
306    )
307    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)