torch_em.data.datasets.light_microscopy.evican

  1"""The EVICAN dataset contains phase contrast and brightfield microscopy images
  2with cell and nucleus segmentations.
  3
  4NOTE: The data is sparsely annotated.
  5
  6The dataset is located at https://edmond.mpg.de/dataset.xhtml?persistentId=doi:10.17617/3.AJBV1S.
  7This dataset is from the publication https://doi.org/10.1093/bioinformatics/btaa225.
  8Please cite it if you use this dataset in your research.
  9"""
 10
 11import os
 12import warnings
 13from glob import glob
 14from tqdm import tqdm
 15from pathlib import Path
 16from natsort import natsorted
 17from typing import Union, Literal, Tuple, Optional, List
 18
 19import numpy as np
 20import imageio.v3 as imageio
 21
 22from torch.utils.data import Dataset, DataLoader
 23
 24import torch_em
 25
 26try:
 27    from pycocotools.coco import COCO
 28except ImportError:
 29    COCO = None
 30
 31from .. import util
 32
 33
 34URLS = {
 35    # Image archives
 36    "images_train": "https://edmond.mpg.de/api/access/datafile/102910",
 37    "images_val": "https://edmond.mpg.de/api/access/datafile/102911",
 38    "images_eval": "https://edmond.mpg.de/api/access/datafile/102912",
 39
 40    # EVICAN2 annotations
 41    "annotations_evican2_train": "https://edmond.mpg.de/api/access/datafile/102915",
 42    "annotations_evican2_val": "https://edmond.mpg.de/api/access/datafile/102916",
 43    "annotations_evican2_eval_easy": "https://edmond.mpg.de/api/access/datafile/102918",
 44    "annotations_evican2_eval_medium": "https://edmond.mpg.de/api/access/datafile/102919",
 45    "annotations_evican2_eval_difficult": "https://edmond.mpg.de/api/access/datafile/102917",
 46
 47    # EVICAN60 annotations
 48    "annotations_evican60_train": "https://edmond.mpg.de/api/access/datafile/102921",
 49    "annotations_evican60_val": "https://edmond.mpg.de/api/access/datafile/102922",
 50    "annotations_evican60_eval_easy": "https://edmond.mpg.de/api/access/datafile/102924",
 51    "annotations_evican60_eval_medium": "https://edmond.mpg.de/api/access/datafile/102920",
 52    "annotations_evican60_eval_difficult": "https://edmond.mpg.de/api/access/datafile/102923",
 53}
 54CHECKSUMS = None
 55
 56ANNOTATION_TYPES = ["evican2", "evican60"]
 57SEGMENTATION_TYPES = ["cell", "nucleus"]
 58SPLITS = ["train", "val", "eval_easy", "eval_medium", "eval_difficult"]
 59
 60# Map from segmentation type to COCO category name.
 61_CATEGORY_NAMES = {"cell": "Cell", "nucleus": "Nucleus"}
 62
 63
 64def _annotations_to_instances(coco, image_metadata, category_ids):
 65    """Convert COCO annotations to an instance segmentation mask.
 66
 67    NOTE: The EVICAN annotations lack the 'area' field, so we compute it from the masks.
 68    """
 69    import vigra
 70
 71    annotation_ids = coco.getAnnIds(imgIds=image_metadata["id"], catIds=category_ids)
 72    annotations = coco.loadAnns(annotation_ids)
 73    assert len(annotations) <= np.iinfo("uint16").max
 74
 75    shape = (image_metadata["height"], image_metadata["width"])
 76    seg = np.zeros(shape, dtype="uint32")
 77
 78    # Compute masks and their areas for sorting.
 79    masks = []
 80    areas = []
 81    for annotation in annotations:
 82        mask = coco.annToMask(annotation).astype("bool")
 83        masks.append(mask)
 84        areas.append(mask.sum())
 85
 86    # Sort by area (smallest last so they are not occluded by larger objects).
 87    sorting = np.argsort(areas)[::-1]
 88
 89    for seg_id, idx in enumerate(sorting, 1):
 90        seg[masks[idx]] = seg_id
 91
 92    # Filter small noise objects from overlapping annotations.
 93    min_size = 50
 94    seg_ids, sizes = np.unique(seg, return_counts=True)
 95    seg[np.isin(seg, seg_ids[sizes < min_size])] = 0
 96
 97    vigra.analysis.relabelConsecutive(seg, out=seg)
 98
 99    return seg.astype("uint16")
100
101
102def _download_images(path, split, download):
103    """Download and extract image archives for the given split."""
104    # eval_easy, eval_medium, eval_difficult all share the same eval images.
105    image_split = "eval" if split.startswith("eval") else split
106    image_dir = os.path.join(path, "images", image_split)
107    if os.path.exists(image_dir):
108        return image_dir
109
110    os.makedirs(image_dir, exist_ok=True)
111    url_key = f"images_{image_split}"
112    zip_path = os.path.join(path, f"EVICAN_{image_split}.zip")
113    util.download_source(zip_path, URLS[url_key], download, checksum=None)
114    util.unzip(zip_path, image_dir, remove=True)
115
116    return image_dir
117
118
119def _download_annotations(path, split, annotation_type, download):
120    """Download the COCO annotation JSON for the given split and annotation type."""
121    ann_dir = os.path.join(path, "annotations")
122    os.makedirs(ann_dir, exist_ok=True)
123
124    ann_file = os.path.join(ann_dir, f"instances_{split}_{annotation_type}.json")
125    if os.path.exists(ann_file):
126        return ann_file
127
128    url_key = f"annotations_{annotation_type}_{split}"
129    util.download_source(ann_file, URLS[url_key], download, checksum=None)
130
131    return ann_file
132
133
134def _create_segmentations_from_coco_annotations(path, split, annotation_type, segmentation_type):
135    """Convert COCO annotations to instance segmentation masks."""
136    assert COCO is not None, (
137        "'pycocotools' is required for processing the EVICAN ground-truth. "
138        "Install it with 'conda install -c conda-forge pycocotools'."
139    )
140
141    image_split = "eval" if split.startswith("eval") else split
142    image_dir = os.path.join(path, "images", image_split)
143    seg_dir = os.path.join(path, "segmentations", annotation_type, segmentation_type, split)
144
145    if os.path.exists(seg_dir):
146        # Check that segmentation files exist.
147        seg_paths = glob(os.path.join(seg_dir, "*.tif"))
148        if len(seg_paths) > 0:
149            image_paths = [os.path.join(image_dir, f"{Path(sp).stem}.jpg") for sp in seg_paths]
150            return natsorted(image_paths), natsorted(seg_paths)
151
152    os.makedirs(seg_dir, exist_ok=True)
153
154    ann_file = os.path.join(path, "annotations", f"instances_{split}_{annotation_type}.json")
155    assert os.path.exists(ann_file), f"Annotation file not found: {ann_file}"
156
157    coco = COCO(ann_file)
158    category_name = _CATEGORY_NAMES[segmentation_type]
159    category_ids = coco.getCatIds(catNms=[category_name])
160    image_ids = coco.getImgIds(catIds=category_ids)
161
162    image_paths, seg_paths = [], []
163    for image_id in tqdm(
164        image_ids, desc=f"Creating EVICAN {segmentation_type} segmentations from COCO annotations ({split})"
165    ):
166        image_metadata = coco.loadImgs(image_id)[0]
167        file_name = image_metadata["file_name"]
168
169        image_path = os.path.join(image_dir, file_name)
170        assert os.path.exists(image_path), image_path
171        image_paths.append(image_path)
172
173        seg_path = os.path.join(seg_dir, f"{Path(file_name).stem}.tif")
174        seg_paths.append(seg_path)
175        if os.path.exists(seg_path):
176            continue
177
178        # Suppress numpy 2.0 deprecation warning from pycocotools.
179        with warnings.catch_warnings():
180            warnings.filterwarnings("ignore", category=DeprecationWarning, module="pycocotools")
181            seg = _annotations_to_instances(coco, image_metadata, category_ids)
182
183        imageio.imwrite(seg_path, seg, compression="zlib")
184
185    assert len(image_paths) == len(seg_paths) and len(image_paths) > 0
186    return natsorted(image_paths), natsorted(seg_paths)
187
188
189def get_evican_data(
190    path: Union[os.PathLike, str],
191    split: Literal["train", "val", "eval_easy", "eval_medium", "eval_difficult"],
192    annotation_type: Literal["evican2", "evican60"] = "evican2",
193    download: bool = False,
194) -> str:
195    """Download the EVICAN dataset.
196
197    Args:
198        path: Filepath to a folder where the downloaded data will be saved.
199        split: The data split to use. One of 'train', 'val', 'eval_easy', 'eval_medium' or 'eval_difficult'.
200        annotation_type: The annotation type to use. Either 'evican2' or 'evican60'.
201        download: Whether to download the data if it is not present.
202
203    Returns:
204        The path where the dataset is stored.
205    """
206    os.makedirs(path, exist_ok=True)
207    _download_images(path, split, download)
208    _download_annotations(path, split, annotation_type, download)
209    return path
210
211
212def get_evican_paths(
213    path: Union[os.PathLike, str],
214    split: Literal["train", "val", "eval_easy", "eval_medium", "eval_difficult"],
215    annotation_type: Literal["evican2", "evican60"] = "evican2",
216    segmentation_type: Literal["cell", "nucleus"] = "cell",
217    download: bool = False,
218) -> Tuple[List[str], List[str]]:
219    """Get paths to the EVICAN data.
220
221    Args:
222        path: Filepath to a folder where the downloaded data will be saved.
223        split: The data split to use. One of 'train', 'val', 'eval_easy', 'eval_medium' or 'eval_difficult'.
224        annotation_type: The annotation type to use. Either 'evican2' or 'evican60'.
225        segmentation_type: The segmentation target. Either 'cell' or 'nucleus'.
226        download: Whether to download the data if it is not present.
227
228    Returns:
229        List of filepaths for the image data.
230        List of filepaths for the label data.
231    """
232    get_evican_data(path, split, annotation_type, download)
233    image_paths, seg_paths = _create_segmentations_from_coco_annotations(
234        path, split, annotation_type, segmentation_type
235    )
236    return image_paths, seg_paths
237
238
239def get_evican_dataset(
240    path: Union[os.PathLike, str],
241    patch_shape: Tuple[int, int],
242    split: Literal["train", "val", "eval_easy", "eval_medium", "eval_difficult"],
243    annotation_type: Literal["evican2", "evican60"] = "evican2",
244    segmentation_type: Literal["cell", "nucleus"] = "cell",
245    offsets: Optional[List[List[int]]] = None,
246    boundaries: bool = False,
247    binary: bool = False,
248    download: bool = False,
249    **kwargs
250) -> Dataset:
251    """Get the EVICAN dataset for cell and nucleus segmentation.
252
253    Args:
254        path: Filepath to a folder where the downloaded data will be saved.
255        patch_shape: The patch shape to use for training.
256        split: The data split to use. One of 'train', 'val', 'eval_easy', 'eval_medium' or 'eval_difficult'.
257        annotation_type: The annotation type to use. Either 'evican2' or 'evican60'.
258        segmentation_type: The segmentation target. Either 'cell' or 'nucleus'.
259        offsets: Offset values for affinity computation used as target.
260        boundaries: Whether to compute boundaries as the target.
261        binary: Whether to use a binary segmentation target.
262        download: Whether to download the data if it is not present.
263        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
264
265    Returns:
266        The segmentation dataset.
267    """
268    assert split in SPLITS, f"'{split}' is not a valid split. Choose from {SPLITS}."
269    assert annotation_type in ANNOTATION_TYPES, f"'{annotation_type}' is not valid. Choose from {ANNOTATION_TYPES}."
270    assert segmentation_type in SEGMENTATION_TYPES, \
271        f"'{segmentation_type}' is not valid. Choose from {SEGMENTATION_TYPES}."
272
273    image_paths, seg_paths = get_evican_paths(path, split, annotation_type, segmentation_type, download)
274
275    kwargs = util.ensure_transforms(ndim=2, **kwargs)
276    kwargs, _ = util.add_instance_label_transform(
277        kwargs, add_binary_target=True, offsets=offsets, boundaries=boundaries, binary=binary
278    )
279
280    return torch_em.default_segmentation_dataset(
281        raw_paths=image_paths,
282        raw_key=None,
283        label_paths=seg_paths,
284        label_key=None,
285        patch_shape=patch_shape,
286        is_seg_dataset=False,
287        **kwargs
288    )
289
290
291def get_evican_loader(
292    path: Union[os.PathLike, str],
293    batch_size: int,
294    patch_shape: Tuple[int, int],
295    split: Literal["train", "val", "eval_easy", "eval_medium", "eval_difficult"],
296    annotation_type: Literal["evican2", "evican60"] = "evican2",
297    segmentation_type: Literal["cell", "nucleus"] = "cell",
298    offsets: Optional[List[List[int]]] = None,
299    boundaries: bool = False,
300    binary: bool = False,
301    download: bool = False,
302    **kwargs
303) -> DataLoader:
304    """Get the EVICAN dataloader for cell and nucleus segmentation.
305
306    Args:
307        path: Filepath to a folder where the downloaded data will be saved.
308        batch_size: The batch size for training.
309        patch_shape: The patch shape to use for training.
310        split: The data split to use. One of 'train', 'val', 'eval_easy', 'eval_medium' or 'eval_difficult'.
311        annotation_type: The annotation type to use. Either 'evican2' or 'evican60'.
312        segmentation_type: The segmentation target. Either 'cell' or 'nucleus'.
313        offsets: Offset values for affinity computation used as target.
314        boundaries: Whether to compute boundaries as the target.
315        binary: Whether to use a binary segmentation target.
316        download: Whether to download the data if it is not present.
317        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
318
319    Returns:
320        The DataLoader.
321    """
322    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
323    dataset = get_evican_dataset(
324        path=path,
325        patch_shape=patch_shape,
326        split=split,
327        annotation_type=annotation_type,
328        segmentation_type=segmentation_type,
329        offsets=offsets,
330        boundaries=boundaries,
331        binary=binary,
332        download=download,
333        **ds_kwargs,
334    )
335    return torch_em.get_data_loader(dataset=dataset, batch_size=batch_size, **loader_kwargs)