torch_em.data.datasets.medical.insect_anatomy

  1"""The InsectAnatomy dataset contains annotations for brain segmentation in micro-CT scans of ant heads.
  2
  3It comprises 2d cross-sections (along the xy, xz and yz planes) of the head scans of 76 ant species with
  4semi-manually segmented brain masks, split into an official training (60%) and testing (40%) set.
  5The images and masks are stored as tif files. The label ids are: 0 = background, 1 = brain.
  6This is the 'InsectAnatomy' out-of-distribution set of the nnInteractive benchmark
  7(https://arxiv.org/abs/2503.08373), which reports 84 micro-CT volumes for it.
  8
  9NOTE: The data is hosted on Dryad (https://doi.org/10.5061/dryad.qz612jmgv), which protects its downloads
 10with a browser check that automated downloads cannot pass. If the automatic download fails, please download
 11'training.zip' and 'testing.zip' from the Dryad page in your browser and place them at '<path>'.
 12
 13NOTE: The layout inside the archives is not documented by the authors. The images and masks are matched by
 14their file names (see `MASK_KEYWORDS`); adjust `get_insect_anatomy_paths` if the layout differs. The archives
 15mix three naming conventions and also contain unpaired files, i.e. species for which only the images or only
 16the masks were deposited, so those are skipped: this yields 35422 pairs for 'train' and 7366 for 'test'.
 17
 18NOTE: The masks are 8 bit images with antialiased borders, where the brain is 255 and the background is 0.
 19They are binarized to the label ids above by the dataset, so a custom 'label_transform' overrides this.
 20
 21NOTE: The masks of the test split are inconsistent and it is therefore not recommended for training. In the
 22training split the brain is marked correctly (247 of 250 sampled masks cover the brighter brain tissue, with a
 23mean foreground of 22%), but in the test split most masks have an inverted polarity, i.e. they mark the
 24background instead of the brain: of 179 sampled masks that cover more than half of the image, only 6 cover the
 25brighter tissue. The deposit does not document this, and the polarity cannot be told apart from a correctly
 26annotated slice of a cropped scan, so the masks are passed through as they are stored and a warning is raised.
 27
 28The dataset is located at https://doi.org/10.5061/dryad.qz612jmgv.
 29
 30This dataset is from the publication https://doi.org/10.1002/ntls.20230010.
 31Please cite it if you use this dataset in your research.
 32"""
 33
 34import os
 35import re
 36import zipfile
 37import warnings
 38from glob import glob
 39from natsort import natsorted
 40from typing import Union, Tuple, Literal, List
 41
 42import numpy as np
 43
 44from torch.utils.data import Dataset, DataLoader
 45
 46import torch_em
 47
 48from .. import util
 49
 50
 51URLS = {
 52    "train": "https://datadryad.org/downloads/file_stream/2625340",
 53    "test": "https://datadryad.org/downloads/file_stream/2625339",
 54}
 55
 56ZIP_NAMES = {"train": "training.zip", "test": "testing.zip"}
 57
 58LABEL_IDS = {"background": 0, "brain": 1}
 59
 60MASK_KEYWORDS = ("mask", "label", "seg", "gt")
 61
 62MANUAL_DOWNLOAD_MSG = (
 63    "The automatic download of '{}' from Dryad failed. Dryad protects its downloads with a browser check. "
 64    "Please download the file manually from https://doi.org/10.5061/dryad.qz612jmgv and place it at '{}'."
 65)
 66
 67
 68def _binarize_labels(labels):
 69    # The masks are 8 bit images with antialiased borders, where the brain is 255 and the background is 0.
 70    return (np.asarray(labels) > 127).astype("uint8")
 71
 72
 73def get_insect_anatomy_data(
 74    path: Union[os.PathLike, str], split: Literal["train", "test"], download: bool = False
 75) -> str:
 76    """Download the InsectAnatomy dataset.
 77
 78    Args:
 79        path: Filepath to a folder where the data is downloaded for further processing.
 80        split: The choice of data split.
 81        download: Whether to download the data if it is not present.
 82
 83    Returns:
 84        Filepath where the data is downloaded.
 85    """
 86    if split not in URLS:
 87        raise ValueError(f"'{split}' is not a valid split.")
 88
 89    data_dir = os.path.join(path, split)
 90    if os.path.exists(data_dir):
 91        return data_dir
 92
 93    os.makedirs(path, exist_ok=True)
 94
 95    zip_path = os.path.join(path, ZIP_NAMES[split])
 96    if not os.path.exists(zip_path):
 97        if not download:
 98            raise RuntimeError(f"Cannot find the data at {zip_path}, but download was set to False.")
 99
100        try:
101            util.download_source(path=zip_path, url=URLS[split], download=download)
102        except Exception:
103            raise RuntimeError(MANUAL_DOWNLOAD_MSG.format(ZIP_NAMES[split], path))
104
105    if not zipfile.is_zipfile(zip_path):
106        os.remove(zip_path)
107        raise RuntimeError(MANUAL_DOWNLOAD_MSG.format(ZIP_NAMES[split], path))
108
109    util.unzip(zip_path=zip_path, dst=data_dir)
110    return data_dir
111
112
113def get_insect_anatomy_paths(
114    path: Union[os.PathLike, str], split: Literal["train", "test"], download: bool = False
115) -> Tuple[List[str], List[str]]:
116    """Get paths to the InsectAnatomy data.
117
118    Args:
119        path: Filepath to a folder where the data is downloaded for further processing.
120        split: The choice of data split.
121        download: Whether to download the data if it is not present.
122
123    Returns:
124        List of filepaths for the image data.
125        List of filepaths for the label data.
126    """
127    data_dir = get_insect_anatomy_data(path, split, download)
128
129    if split == "test":
130        warnings.warn(
131            "The masks of the InsectAnatomy test split are inconsistent: most of them mark the background "
132            "instead of the brain. See 'torch_em.data.datasets.medical.insect_anatomy' for details."
133        )
134
135    # The images and their masks are stored as tif files, the masks are marked by 'mask' in the file name.
136    # NOTE: The archives use three different naming conventions, so the name of the image is derived from the
137    # name of the mask by removing the first mask marker: '<name>.tif' pairs with '<name>_mask.tif', and
138    # '<species><slice>.tif' pairs with both '<species>__mask<slice>.tif' and '<species>_mask<slice>.tif'.
139    all_paths = natsorted(glob(os.path.join(data_dir, "**", "*.tif*"), recursive=True))
140
141    # The archives contain the sidecar files that macOS adds when zipping, which are not image data.
142    all_paths = [p for p in all_paths if "__MACOSX" not in p and not os.path.basename(p).startswith("._")]
143
144    mask_paths, images_by_name = [], {}
145    for filepath in all_paths:
146        name = os.path.splitext(os.path.basename(filepath))[0]
147        if any(kw in name.lower() for kw in MASK_KEYWORDS):
148            mask_paths.append(filepath)
149        else:
150            images_by_name[name] = filepath
151
152    raw_paths, label_paths = [], []
153    for filepath in mask_paths:
154        name = os.path.splitext(os.path.basename(filepath))[0]
155        image_name = re.sub(r"_{0,2}mask", "", name, count=1)
156        if image_name in images_by_name:
157            raw_paths.append(images_by_name[image_name])
158            label_paths.append(filepath)
159
160    if len(raw_paths) == 0:
161        raise RuntimeError(
162            f"Found {len(images_by_name)} images and {len(mask_paths)} masks in '{data_dir}', but could not match "
163            "any of them. The images and masks are matched by their file names, see 'MASK_KEYWORDS'."
164        )
165
166    return raw_paths, label_paths
167
168
169def get_insect_anatomy_dataset(
170    path: Union[os.PathLike, str],
171    patch_shape: Tuple[int, int],
172    split: Literal["train", "test"],
173    resize_inputs: bool = False,
174    download: bool = False,
175    **kwargs
176) -> Dataset:
177    """Get the InsectAnatomy dataset for ant brain segmentation.
178
179    Args:
180        path: Filepath to a folder where the data is downloaded for further processing.
181        patch_shape: The patch shape to use for training.
182        split: The choice of data split.
183        resize_inputs: Whether to resize inputs to the desired patch shape.
184        download: Whether to download the data if it is not present.
185        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
186
187    Returns:
188        The segmentation dataset.
189    """
190    raw_paths, label_paths = get_insect_anatomy_paths(path, split, download)
191
192    # The masks are stored as 8 bit images with antialiased borders, so they are binarized to the documented ids.
193    kwargs = util.update_kwargs(kwargs, "label_transform", _binarize_labels)
194
195    if resize_inputs:
196        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
197        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
198            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
199        )
200
201    return torch_em.default_segmentation_dataset(
202        raw_paths=raw_paths,
203        raw_key=None,
204        label_paths=label_paths,
205        label_key=None,
206        patch_shape=patch_shape,
207        is_seg_dataset=False,
208        **kwargs
209    )
210
211
212def get_insect_anatomy_loader(
213    path: Union[os.PathLike, str],
214    batch_size: int,
215    patch_shape: Tuple[int, int],
216    split: Literal["train", "test"],
217    resize_inputs: bool = False,
218    download: bool = False,
219    **kwargs
220) -> DataLoader:
221    """Get the InsectAnatomy dataloader for ant brain segmentation.
222
223    Args:
224        path: Filepath to a folder where the data is downloaded for further processing.
225        batch_size: The batch size for training.
226        patch_shape: The patch shape to use for training.
227        split: The choice of data split.
228        resize_inputs: Whether to resize inputs to the desired patch shape.
229        download: Whether to download the data if it is not present.
230        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
231
232    Returns:
233        The DataLoader.
234    """
235    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
236    dataset = get_insect_anatomy_dataset(path, patch_shape, split, resize_inputs, download, **ds_kwargs)
237    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)