torch_em.data.datasets.histopathology.dcsa_net

  1"""The DCSA-Net dataset contains binary nucleus segmentation masks for H&E stained
  2prostate cancer histopathology images from two internal cohorts (RUMC, YUHS).
  3
  4NOTE: The masks are binary nucleus foreground/background maps, not per-nucleus instance
  5labels. Visual inspection also suggests the annotations are not exhaustive: some visible
  6nuclei in the raw images have no corresponding mask region.
  7
  8This dataset is located at https://doi.org/10.6084/m9.figshare.22249291.
  9This dataset is from the publication https://doi.org/10.3389/fonc.2023.1009681.
 10Please cite it if you use this dataset for your research.
 11"""
 12
 13import os
 14from natsort import natsorted
 15from typing import Union, Literal, Tuple, List
 16
 17import json
 18import pandas as pd
 19import imageio.v3 as imageio
 20from sklearn.model_selection import train_test_split
 21
 22from torch.utils.data import Dataset, DataLoader
 23
 24import torch_em
 25
 26from .. import util
 27
 28
 29URL = "https://ndownloader.figshare.com/files/39539971"
 30CHECKSUM = "25bb4a37672809c7f762a20929218cef838e94c1c316ad1e9c31d801447197ad"
 31
 32
 33def get_dcsa_net_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 34    """Download the DCSA-Net data.
 35
 36    Args:
 37        path: Filepath to a folder where the downloaded data will be saved.
 38        download: Whether to download the data if it is not present.
 39
 40    Returns:
 41        Filepath where the dataset is downloaded and stored for further preprocessing.
 42    """
 43    data_dir = os.path.join(path, "Training Data")
 44    if os.path.exists(data_dir):
 45        return data_dir
 46
 47    os.makedirs(path, exist_ok=True)
 48    zip_path = os.path.join(path, "dcsa_net.zip")
 49    util.download_source(path=zip_path, url=URL, download=download, checksum=CHECKSUM)
 50    util.unzip(zip_path=zip_path, dst=path)
 51
 52    return data_dir
 53
 54
 55def _raw_mask_pairs(data_dir):
 56    image_dir = os.path.join(data_dir, "Train_Images")
 57    mask_dir = os.path.join(data_dir, "Train_Masks")
 58
 59    # The archive stores the RUMC cohort as '.bmp' (indices 1-45) and the YUHS cohort as
 60    # '.jpg' (indices 1-30) under a shared 'Train_Images' folder, while all 75 masks live
 61    # in a single flat 'Train_Masks' sequence: RUMC masks 1-45, then YUHS masks 46-75.
 62    pairs = []
 63    for index in range(1, 46):
 64        pairs.append((
 65            os.path.join(image_dir, f"Prostate ({index}).bmp"),
 66            os.path.join(mask_dir, f"Prostate ({index}).bmp"),
 67        ))
 68    for index in range(1, 31):
 69        pairs.append((
 70            os.path.join(image_dir, f"Prostate ({index}).jpg"),
 71            os.path.join(mask_dir, f"Prostate ({index + 45}).bmp"),
 72        ))
 73
 74    for raw_path, mask_path in pairs:
 75        assert os.path.exists(raw_path), raw_path
 76        assert os.path.exists(mask_path), mask_path
 77
 78    return pairs
 79
 80
 81def _convert_masks_to_binary(data_dir, pairs):
 82    converted_dir = os.path.join(data_dir, "Train_Masks_binary")
 83    os.makedirs(converted_dir, exist_ok=True)
 84
 85    label_paths = []
 86    for _, mask_path in pairs:
 87        out_path = os.path.join(converted_dir, os.path.basename(mask_path))
 88        if not os.path.exists(out_path):
 89            mask = imageio.imread(mask_path)
 90            if mask.ndim == 3:
 91                mask = mask[..., 0]
 92            imageio.imwrite(out_path, (mask > 127).astype("uint8"))
 93        label_paths.append(out_path)
 94
 95    return label_paths
 96
 97
 98def _create_split_csv(path, raw_paths):
 99    csv_path = os.path.join(path, "dcsa_net_split.csv")
100    if os.path.exists(csv_path):
101        df = pd.read_csv(csv_path)
102        return {split: json.loads(df.iloc[0][split].replace("'", '"')) for split in ("train", "val", "test")}
103
104    print(f"Creating a new split file at '{csv_path}'.")
105    image_ids = natsorted(os.path.basename(p) for p in raw_paths)
106
107    train_ids, test_ids = train_test_split(image_ids, test_size=0.2, random_state=42)
108    train_ids, val_ids = train_test_split(train_ids, test_size=0.15, random_state=42)
109    split_ids = {"train": train_ids, "val": val_ids, "test": test_ids}
110
111    df = pd.DataFrame.from_dict([split_ids])
112    df.to_csv(csv_path, index=False)
113
114    return split_ids
115
116
117def get_dcsa_net_paths(
118    path: Union[os.PathLike, str], split: Literal["train", "val", "test"], download: bool = False,
119) -> Tuple[List[str], List[str]]:
120    """Get paths to the DCSA-Net data.
121
122    NOTE: The source publishes no official split, so this function creates and stores a
123    deterministic split (65% train, 15% val, 20% test) the first time it is called.
124
125    Args:
126        path: Filepath to a folder where the downloaded data will be saved.
127        split: The choice of data split.
128        download: Whether to download the data if it is not present.
129
130    Returns:
131        List of filepaths to the image data.
132        List of filepaths to the label data.
133    """
134    data_dir = get_dcsa_net_data(path, download)
135    pairs = _raw_mask_pairs(data_dir)
136    converted_label_paths = _convert_masks_to_binary(data_dir, pairs)
137    raw_paths = [raw_path for raw_path, _ in pairs]
138
139    split_ids = _create_split_csv(path, raw_paths)[split]
140    kept = natsorted(
141        (raw_path, label_path) for raw_path, label_path in zip(raw_paths, converted_label_paths)
142        if os.path.basename(raw_path) in split_ids
143    )
144    raw_paths = [raw_path for raw_path, _ in kept]
145    label_paths = [label_path for _, label_path in kept]
146
147    assert len(raw_paths) == len(label_paths) and len(raw_paths) > 0
148    return raw_paths, label_paths
149
150
151def get_dcsa_net_dataset(
152    path: Union[os.PathLike, str],
153    patch_shape: Tuple[int, int],
154    split: Literal["train", "val", "test"],
155    resize_inputs: bool = False,
156    download: bool = False,
157    **kwargs
158) -> Dataset:
159    """Get the DCSA-Net dataset for nucleus segmentation.
160
161    Args:
162        path: Filepath to a folder where the downloaded data will be saved.
163        patch_shape: The patch shape to use for training.
164        split: The choice of data split.
165        resize_inputs: Whether to resize the inputs.
166        download: Whether to download the data if it is not present.
167        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
168
169    Returns:
170        The segmentation dataset.
171    """
172    raw_paths, label_paths = get_dcsa_net_paths(path, split, download)
173
174    if resize_inputs:
175        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": True}
176        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
177            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
178        )
179
180    return torch_em.default_segmentation_dataset(
181        raw_paths=raw_paths,
182        raw_key=None,
183        label_paths=label_paths,
184        label_key=None,
185        is_seg_dataset=False,
186        patch_shape=patch_shape,
187        with_channels=True,
188        ndim=2,
189        **kwargs
190    )
191
192
193def get_dcsa_net_loader(
194    path: Union[os.PathLike, str],
195    batch_size: int,
196    patch_shape: Tuple[int, int],
197    split: Literal["train", "val", "test"],
198    resize_inputs: bool = False,
199    download: bool = False,
200    **kwargs
201) -> DataLoader:
202    """Get the DCSA-Net dataloader for nucleus segmentation.
203
204    Args:
205        path: Filepath to a folder where the downloaded data will be saved.
206        batch_size: The batch size for training.
207        patch_shape: The patch shape to use for training.
208        split: The choice of data split.
209        resize_inputs: Whether to resize the inputs.
210        download: Whether to download the data if it is not present.
211        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
212
213    Returns:
214        The DataLoader.
215    """
216    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
217    dataset = get_dcsa_net_dataset(path, patch_shape, split, resize_inputs, download, **ds_kwargs)
218    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)