torch_em.data.datasets.medical.particleseg3d

  1"""The ParticleSeg3D dataset contains annotations for individual particle segmentation in micro CT images
  2of mineral, ore, slag and recycling samples.
  3
  4The data consists of 54 annotated micro CT patches (41 for training and 13 for testing) that were extracted from
  528 scanned samples of eight different materials. Each patch is paired with an instance segmentation in which every
  6individual particle carries its own id (0 is the background), so the label ids are instance ids and not semantic ids.
  7The 'split' argument exposes the official train / test split of the publication.
  8
  9The data is hosted in a public share of the DESY sync-and-share service. This module downloads the 'Patches' folder
 10of that share (ca. 2.5 GB), which contains the annotated image patches ('images') and the corresponding instance
 11segmentations ('instance_seg') as nifti files. The 'Samples' folder of the share holds the full micro CT scans
 12(ca. 70 GB) without annotations and is therefore not downloaded here.
 13
 14The dataset is located at https://syncandshare.desy.de/index.php/s/wjiDQ49KangiPj5 and the code of the publication
 15at https://github.com/MIC-DKFZ/ParticleSeg3D.
 16
 17This dataset is from the publication https://doi.org/10.1016/j.powtec.2023.119286.
 18Please cite it if you use this dataset in your research.
 19"""
 20
 21import os
 22import re
 23from glob import glob
 24from tqdm import tqdm
 25from natsort import natsorted
 26from urllib.parse import quote, unquote
 27from typing import Union, Tuple, List, Literal
 28
 29import requests
 30
 31from torch.utils.data import Dataset, DataLoader
 32
 33import torch_em
 34
 35from .. import util
 36
 37
 38SHARE_TOKEN = "wjiDQ49KangiPj5"
 39
 40URL = f"https://syncandshare.desy.de/index.php/s/{SHARE_TOKEN}"
 41
 42WEBDAV_URL = "https://syncandshare.desy.de/public.php/webdav"
 43
 44# The files are downloaded individually from the public share, so there is no checksum for a single archive.
 45CHECKSUM = None
 46
 47# The number of annotated patches per split.
 48N_PATCHES = {"train": 41, "test": 13}
 49
 50FOLDERS = ["images", "instance_seg"]
 51
 52
 53def _list_share_folder(folder):
 54    """List the file names in a folder of the public share via the WebDAV endpoint of Nextcloud.
 55
 56    Public shares are accessed by using the share token as the user name and an empty password.
 57    """
 58    response = requests.request(
 59        "PROPFIND", f"{WEBDAV_URL}/{quote(folder)}/", headers={"Depth": "1"}, auth=(SHARE_TOKEN, "")
 60    )
 61    response.raise_for_status()
 62
 63    fnames = []
 64    for href in re.findall(r"<d:href>(.*?)</d:href>", response.text):
 65        fname = unquote(href).rstrip("/").split("/")[-1]
 66        if fname.endswith(".nii.gz"):
 67            fnames.append(fname)
 68    return natsorted(fnames)
 69
 70
 71def _download_share_folder(folder, dst, download):
 72    """Download all nifti files of a folder of the public share into `dst`."""
 73    if os.path.exists(dst):
 74        return
 75
 76    if not download:
 77        raise RuntimeError(f"Cannot find the data at {dst}, but download was set to False.")
 78
 79    tmp_dir = f"{dst}.tmp"
 80    os.makedirs(tmp_dir, exist_ok=True)
 81    fnames = _list_share_folder(folder)
 82    for fname in tqdm(fnames, desc=f"Download {len(fnames)} files from '{folder}'"):
 83        out_path = os.path.join(tmp_dir, fname)
 84        if os.path.exists(out_path):
 85            continue
 86        url = f"{URL}/download?path={quote('/' + folder)}&files={quote(fname)}"
 87        util.download_source(path=out_path, url=url, download=download, checksum=CHECKSUM)
 88
 89    os.rename(tmp_dir, dst)
 90
 91
 92def get_particleseg3d_data(
 93    path: Union[os.PathLike, str], split: Literal["train", "test"], download: bool = False
 94) -> str:
 95    """Download the ParticleSeg3D dataset.
 96
 97    Args:
 98        path: Filepath to a folder where the data is downloaded for further processing.
 99        split: The choice of data split. Either 'train' or 'test'.
100        download: Whether to download the data if it is not present.
101
102    Returns:
103        Filepath where the data is downloaded.
104    """
105    if split not in N_PATCHES:
106        raise ValueError(f"'{split}' is not a valid split. Please choose one of {list(N_PATCHES.keys())}.")
107
108    data_dir = os.path.join(path, "Patches", split)
109    os.makedirs(path, exist_ok=True)
110    for folder in FOLDERS:
111        _download_share_folder(f"Patches/{split}/{folder}", os.path.join(data_dir, folder), download)
112
113    return data_dir
114
115
116def get_particleseg3d_paths(
117    path: Union[os.PathLike, str], split: Literal["train", "test"], download: bool = False
118) -> Tuple[List[str], List[str]]:
119    """Get paths to the ParticleSeg3D data.
120
121    Args:
122        path: Filepath to a folder where the data is downloaded for further processing.
123        split: The choice of data split. Either 'train' or 'test'.
124        download: Whether to download the data if it is not present.
125
126    Returns:
127        List of filepaths for the image data.
128        List of filepaths for the label data.
129    """
130    data_dir = get_particleseg3d_data(path, split, download)
131
132    raw_paths = natsorted(glob(os.path.join(data_dir, "images", "*.nii.gz")))
133    label_paths = [p.replace(os.sep + "images" + os.sep, os.sep + "instance_seg" + os.sep) for p in raw_paths]
134    assert len(raw_paths) == N_PATCHES[split] and all(os.path.exists(p) for p in label_paths)
135
136    return raw_paths, label_paths
137
138
139def get_particleseg3d_dataset(
140    path: Union[os.PathLike, str],
141    patch_shape: Tuple[int, ...],
142    split: Literal["train", "test"],
143    resize_inputs: bool = False,
144    download: bool = False,
145    **kwargs
146) -> Dataset:
147    """Get the ParticleSeg3D dataset for particle instance segmentation.
148
149    Args:
150        path: Filepath to a folder where the data is downloaded for further processing.
151        patch_shape: The patch shape to use for training.
152        split: The choice of data split. Either 'train' or 'test'.
153        resize_inputs: Whether to resize inputs to the desired patch shape.
154        download: Whether to download the data if it is not present.
155        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
156
157    Returns:
158        The segmentation dataset.
159    """
160    raw_paths, label_paths = get_particleseg3d_paths(path, split, download)
161
162    if resize_inputs:
163        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
164        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
165            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
166        )
167
168    return torch_em.default_segmentation_dataset(
169        raw_paths=raw_paths,
170        raw_key="data",
171        label_paths=label_paths,
172        label_key="data",
173        patch_shape=patch_shape,
174        is_seg_dataset=True,
175        **kwargs
176    )
177
178
179def get_particleseg3d_loader(
180    path: Union[os.PathLike, str],
181    batch_size: int,
182    patch_shape: Tuple[int, ...],
183    split: Literal["train", "test"],
184    resize_inputs: bool = False,
185    download: bool = False,
186    **kwargs
187) -> DataLoader:
188    """Get the ParticleSeg3D dataloader for particle instance segmentation.
189
190    Args:
191        path: Filepath to a folder where the data is downloaded for further processing.
192        batch_size: The batch size for training.
193        patch_shape: The patch shape to use for training.
194        split: The choice of data split. Either 'train' or 'test'.
195        resize_inputs: Whether to resize inputs to the desired patch shape.
196        download: Whether to download the data if it is not present.
197        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
198
199    Returns:
200        The DataLoader.
201    """
202    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
203    dataset = get_particleseg3d_dataset(path, patch_shape, split, resize_inputs, download, **ds_kwargs)
204    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)