torch_em.data.datasets.light_microscopy.urinary_tract

  1"""The UrinaryTract dataset contains annotations for urinary cells in
  2bright-field microscopy images of unstained and untreated urine.
  3
  4The dataset holds 300 images of voided urine from patients with a symptomatic urinary tract
  5infection. Two experts labelled 3562 cells and assigned each one to one of seven clinically
  6significant classes. Every image comes with a foreground mask and a class mask.
  7
  8NOTE: The dataset provides semantic labels only. The masks store class ids, not instance ids, so
  9this loader cannot return cell instances. Connected components do not recover the cells either,
 10because a single cell often breaks into several components in the masks.
 11
 12NOTE: The three splits hold 100 images each, but they differ a lot in cell density.
 13
 14The dataset is located at https://doi.org/10.14278/rodare.2473 under the CC BY 4.0 license.
 15This dataset is from the publication https://doi.org/10.1038/s41597-024-02975-0.
 16Please cite it if you use this dataset in your research.
 17"""
 18
 19import os
 20from glob import glob
 21from natsort import natsorted
 22from typing import List, Literal, Tuple, Union
 23
 24from torch.utils.data import DataLoader, Dataset
 25
 26import torch_em
 27
 28from .. import util
 29
 30
 31URL = "https://rodare.hzdr.de/record/2473/files/ds1.zip"
 32CHECKSUM = "ae66af80c2c0d589c8fc6be21327988cf3ab2c2ed1ccaeadc5783b6e6dd51f95"
 33
 34SPLITS = ("train", "validation", "test")
 35
 36LABEL_CHOICES = {"binary": "bin_mask", "semantic": "mult_mask"}
 37
 38# The class ids of the multi class masks, see Table 1 of the publication.
 39CLASS_NAMES = {
 40    1: "rod",
 41    2: "rbc_wbc",
 42    3: "yeast",
 43    4: "miscellaneous",
 44    5: "single_epc",
 45    6: "small_epc_sheet",
 46    7: "large_epc_sheet",
 47}
 48
 49
 50def get_urinary_tract_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 51    """Download the UrinaryTract dataset.
 52
 53    Args:
 54        path: Filepath to a folder where the downloaded data will be saved.
 55        download: Whether to download the data if it is not present.
 56
 57    Returns:
 58        The filepath to the extracted data.
 59    """
 60    data_dir = os.path.join(path, "ds1")
 61    if os.path.exists(data_dir):
 62        return data_dir
 63
 64    os.makedirs(path, exist_ok=True)
 65    zip_path = os.path.join(path, "ds1.zip")
 66    util.download_source(zip_path, URL, download, CHECKSUM)
 67    util.unzip(zip_path=zip_path, dst=path)
 68
 69    return data_dir
 70
 71
 72def get_urinary_tract_paths(
 73    path: Union[os.PathLike, str],
 74    split: Literal["train", "validation", "test"] = "train",
 75    label_choice: Literal["binary", "semantic"] = "semantic",
 76    download: bool = False,
 77) -> Tuple[List[str], List[str]]:
 78    """Get paths to the UrinaryTract data.
 79
 80    Args:
 81        path: Filepath to a folder where the downloaded data will be saved.
 82        split: The data split. Either 'train', 'validation' or 'test'.
 83        label_choice: The label to use. Either 'semantic' for the seven classes, or 'binary' for
 84            the foreground mask.
 85        download: Whether to download the data if it is not present.
 86
 87    Returns:
 88        List of filepaths for the image data.
 89        List of filepaths for the label data.
 90    """
 91    if split not in SPLITS:
 92        raise ValueError(f"'{split}' is not a valid split. Choose from {list(SPLITS)}.")
 93    if label_choice not in LABEL_CHOICES:
 94        raise ValueError(f"'{label_choice}' is not a valid label choice. Choose from {list(LABEL_CHOICES)}.")
 95
 96    data_dir = get_urinary_tract_data(path, download)
 97    label_dir = LABEL_CHOICES[label_choice]
 98
 99    image_paths = natsorted(glob(os.path.join(data_dir, split, "img", "cls", "*.tif")))
100    label_paths = natsorted(glob(os.path.join(data_dir, split, label_dir, "cls", "*.tif")))
101
102    if not image_paths:
103        raise RuntimeError(f"Could not find any UrinaryTract images in {data_dir}.")
104    if len(image_paths) != len(label_paths):
105        raise RuntimeError(
106            f"Found {len(image_paths)} images but {len(label_paths)} labels for the '{split}' split."
107        )
108
109    return image_paths, label_paths
110
111
112def get_urinary_tract_dataset(
113    path: Union[os.PathLike, str],
114    patch_shape: Tuple[int, int],
115    split: Literal["train", "validation", "test"] = "train",
116    label_choice: Literal["binary", "semantic"] = "semantic",
117    download: bool = False,
118    **kwargs,
119) -> Dataset:
120    """Get the UrinaryTract dataset for urinary cell segmentation.
121
122    Args:
123        path: Filepath to a folder where the downloaded data will be saved.
124        patch_shape: The 2D patch shape to use for training.
125        split: The data split. Either 'train', 'validation' or 'test'.
126        label_choice: The label to use. Either 'semantic' for the seven classes, or 'binary' for
127            the foreground mask.
128        download: Whether to download the data if it is not present.
129        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
130
131    Returns:
132        The segmentation dataset.
133    """
134    if len(patch_shape) != 2:
135        raise ValueError(f"The UrinaryTract patch shape must be two-dimensional, got {patch_shape}.")
136
137    image_paths, label_paths = get_urinary_tract_paths(path, split, label_choice, download)
138
139    if label_choice == "binary":
140        # The masks store 0 and 255, so map them to a background and a foreground id.
141        kwargs["label_transform"] = torch_em.transform.label.labels_to_binary
142
143    kwargs = util.ensure_transforms(ndim=2, **kwargs)
144
145    return torch_em.default_segmentation_dataset(
146        raw_paths=image_paths,
147        raw_key=None,
148        label_paths=label_paths,
149        label_key=None,
150        patch_shape=patch_shape,
151        is_seg_dataset=False,
152        ndim=2,
153        **kwargs,
154    )
155
156
157def get_urinary_tract_loader(
158    path: Union[os.PathLike, str],
159    batch_size: int,
160    patch_shape: Tuple[int, int],
161    split: Literal["train", "validation", "test"] = "train",
162    label_choice: Literal["binary", "semantic"] = "semantic",
163    download: bool = False,
164    **kwargs,
165) -> DataLoader:
166    """Get the UrinaryTract dataloader for urinary cell segmentation.
167
168    Args:
169        path: Filepath to a folder where the downloaded data will be saved.
170        batch_size: The batch size for training.
171        patch_shape: The 2D patch shape to use for training.
172        split: The data split. Either 'train', 'validation' or 'test'.
173        label_choice: The label to use. Either 'semantic' for the seven classes, or 'binary' for
174            the foreground mask.
175        download: Whether to download the data if it is not present.
176        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or the PyTorch DataLoader.
177
178    Returns:
179        The DataLoader.
180    """
181    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
182    dataset = get_urinary_tract_dataset(
183        path=path,
184        patch_shape=patch_shape,
185        split=split,
186        label_choice=label_choice,
187        download=download,
188        **ds_kwargs,
189    )
190    return torch_em.get_data_loader(dataset, batch_size=batch_size, **loader_kwargs)