torch_em.data.datasets.histopathology.cytonuke

  1"""The CytoNuke dataset contains annotations for nucleus and whole-cell segmentation
  2in H&E stained bright-field histopathology images of head and neck squamous cell
  3carcinoma. The images are re-distributed crops from the CPTAC-HNSCC collection.
  4
  5This dataset is located at https://zenodo.org/records/10560728.
  6This dataset is from the publication https://doi.org/10.1016/j.cmpb.2024.108215.
  7Please cite it if you use this dataset for your research.
  8"""
  9
 10import os
 11from tqdm import tqdm
 12from natsort import natsorted
 13from typing import Union, Literal, Tuple, List
 14
 15import json
 16import numpy as np
 17import pandas as pd
 18import imageio.v3 as imageio
 19from skimage.draw import polygon as sk_polygon
 20from skimage.segmentation import relabel_sequential
 21from sklearn.model_selection import train_test_split
 22
 23from torch.utils.data import Dataset, DataLoader
 24
 25import torch_em
 26
 27from .. import util
 28
 29
 30URL = "https://zenodo.org/records/10560728/files/CytoNuke%20Dataset.zip"
 31CHECKSUM = "834836444fe749cc48daa7557458ba578bc00306a59286896e3d29666e98777a"
 32
 33CATEGORY_IDS = {"nuclei": 0, "cell": 1}
 34
 35
 36def _annotations_to_instances(coco, image_metadata, category_id):
 37    height, width = image_metadata["height"], image_metadata["width"]
 38    seg = np.zeros((height, width), dtype="uint32")
 39
 40    annotations = [a for a in coco["annotations"] if a["image_id"] == image_metadata["id"]]
 41    annotations = [a for a in annotations if a["category_id"] == category_id]
 42
 43    for seg_id, annotation in enumerate(annotations, 1):
 44        for polygon in annotation["segmentation"]:
 45            xs, ys = polygon[0::2], polygon[1::2]
 46            rr, cc = sk_polygon(ys, xs, shape=(height, width))
 47            seg[rr, cc] = seg_id
 48
 49    seg, _, _ = relabel_sequential(seg)
 50    return seg.astype("uint16")
 51
 52
 53def _create_segmentations_from_annotations(data_dir, annotations):
 54    image_dir = os.path.join(data_dir, "images")
 55    seg_dir = os.path.join(data_dir, "labels", annotations)
 56    os.makedirs(seg_dir, exist_ok=True)
 57
 58    with open(os.path.join(data_dir, "coco.json")) as f:
 59        coco = json.load(f)
 60
 61    category_id = CATEGORY_IDS[annotations]
 62
 63    image_paths, seg_paths = [], []
 64    for image_metadata in tqdm(coco["images"], desc=f"Creating '{annotations}' segmentations from coco annotations"):
 65        file_name = image_metadata["file_name"]
 66        image_path = os.path.join(image_dir, file_name)
 67        assert os.path.exists(image_path), image_path
 68        image_paths.append(image_path)
 69
 70        seg_path = os.path.join(seg_dir, file_name.replace(".png", ".tif"))
 71        seg_paths.append(seg_path)
 72        if os.path.exists(seg_path):
 73            continue
 74
 75        seg = _annotations_to_instances(coco, image_metadata, category_id)
 76        imageio.imwrite(seg_path, seg, compression="zlib")
 77
 78    return natsorted(image_paths), natsorted(seg_paths)
 79
 80
 81def _create_split_csv(path, image_paths):
 82    csv_path = os.path.join(path, "cytonuke_split.csv")
 83    if os.path.exists(csv_path):
 84        df = pd.read_csv(csv_path)
 85        return {split: json.loads(df.iloc[0][split].replace("'", '"')) for split in ("train", "val", "test")}
 86
 87    print(f"Creating a new split file at '{csv_path}'.")
 88    image_ids = natsorted(os.path.basename(p) for p in image_paths)
 89
 90    train_ids, test_ids = train_test_split(image_ids, test_size=0.2, random_state=42)
 91    train_ids, val_ids = train_test_split(train_ids, test_size=0.15, random_state=42)
 92    split_ids = {"train": train_ids, "val": val_ids, "test": test_ids}
 93
 94    df = pd.DataFrame.from_dict([split_ids])
 95    df.to_csv(csv_path, index=False)
 96
 97    return split_ids
 98
 99
100def get_cytonuke_data(path: Union[os.PathLike, str], download: bool = False) -> str:
101    """Download the CytoNuke data.
102
103    Args:
104        path: Filepath to a folder where the downloaded data will be saved.
105        download: Whether to download the data if it is not present.
106
107    Returns:
108        Filepath where the dataset is downloaded and stored for further preprocessing.
109    """
110    data_dir = os.path.join(path, "data")
111    if os.path.exists(data_dir):
112        return data_dir
113
114    os.makedirs(path, exist_ok=True)
115    zip_path = os.path.join(path, "cytonuke.zip")
116    util.download_source(path=zip_path, url=URL, download=download, checksum=CHECKSUM)
117    util.unzip(zip_path=zip_path, dst=data_dir)
118
119    return data_dir
120
121
122def get_cytonuke_paths(
123    path: Union[os.PathLike, str],
124    split: Literal["train", "val", "test"],
125    annotations: Literal["nuclei", "cell"] = "cell",
126    download: bool = False,
127) -> Tuple[List[str], List[str]]:
128    """Get paths to the CytoNuke data.
129
130    NOTE: The source publishes no official split, so this function creates and stores a
131    deterministic split (65% train, 15% val, 20% test) the first time it is called.
132
133    Args:
134        path: Filepath to a folder where the downloaded data will be saved.
135        split: The choice of data split.
136        annotations: The choice of annotations.
137        download: Whether to download the data if it is not present.
138
139    Returns:
140        List of filepaths to the image data.
141        List of filepaths to the label data.
142    """
143    if annotations not in CATEGORY_IDS:
144        raise ValueError(f"'{annotations}' is not a valid annotation choice.")
145
146    data_dir = get_cytonuke_data(path, download)
147    image_paths, seg_paths = _create_segmentations_from_annotations(data_dir, annotations)
148
149    split_ids = _create_split_csv(path, image_paths)[split]
150    raw_paths = [p for p in image_paths if os.path.basename(p) in split_ids]
151    label_paths = [p for p in seg_paths if os.path.basename(p).replace(".tif", ".png") in split_ids]
152
153    assert len(raw_paths) == len(label_paths) and len(raw_paths) > 0
154    return raw_paths, label_paths
155
156
157def get_cytonuke_dataset(
158    path: Union[os.PathLike, str],
159    patch_shape: Tuple[int, int],
160    split: Literal["train", "val", "test"],
161    annotations: Literal["nuclei", "cell"] = "cell",
162    resize_inputs: bool = False,
163    download: bool = False,
164    **kwargs
165) -> Dataset:
166    """Get the CytoNuke dataset for nucleus and whole-cell segmentation.
167
168    Args:
169        path: Filepath to a folder where the downloaded data will be saved.
170        patch_shape: The patch shape to use for training.
171        split: The choice of data split.
172        annotations: The choice of annotations.
173        resize_inputs: Whether to resize the inputs.
174        download: Whether to download the data if it is not present.
175        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
176
177    Returns:
178        The segmentation dataset.
179    """
180    raw_paths, label_paths = get_cytonuke_paths(path, split, annotations, download)
181
182    if resize_inputs:
183        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": True}
184        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
185            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
186        )
187
188    return torch_em.default_segmentation_dataset(
189        raw_paths=raw_paths,
190        raw_key=None,
191        label_paths=label_paths,
192        label_key=None,
193        is_seg_dataset=False,
194        patch_shape=patch_shape,
195        with_channels=True,
196        ndim=2,
197        **kwargs
198    )
199
200
201def get_cytonuke_loader(
202    path: Union[os.PathLike, str],
203    batch_size: int,
204    patch_shape: Tuple[int, int],
205    split: Literal["train", "val", "test"],
206    annotations: Literal["nuclei", "cell"] = "cell",
207    resize_inputs: bool = False,
208    download: bool = False,
209    **kwargs
210) -> DataLoader:
211    """Get the CytoNuke dataloader for nucleus and whole-cell segmentation.
212
213    Args:
214        path: Filepath to a folder where the downloaded data will be saved.
215        batch_size: The batch size for training.
216        patch_shape: The patch shape to use for training.
217        split: The choice of data split.
218        annotations: The choice of annotations.
219        resize_inputs: Whether to resize the inputs.
220        download: Whether to download the data if it is not present.
221        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
222
223    Returns:
224        The DataLoader.
225    """
226    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
227    dataset = get_cytonuke_dataset(path, patch_shape, split, annotations, resize_inputs, download, **ds_kwargs)
228    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)