torch_em.data.datasets.histopathology.cryonuseg

  1"""The CryoNuSeg dataset contains annotations for nucleus segmentation
  2in cryosectioned H&E stained histological images of 10 different organs.
  3
  4This dataset is from the publication https://doi.org/10.1016/j.compbiomed.2021.104349.
  5Please cite it if you use this dataset for your research.
  6"""
  7
  8import os
  9from glob import glob
 10from natsort import natsorted
 11from typing import Union, Tuple, Literal, List
 12
 13import json
 14import pandas as pd
 15from sklearn.model_selection import train_test_split
 16
 17from torch.utils.data import Dataset, DataLoader
 18
 19import torch_em
 20
 21from .. import util
 22
 23
 24def _create_split_csv(path, data_dir, split):
 25    csv_path = os.path.join(path, 'cryonuseg_split.csv')
 26    if os.path.exists(csv_path):
 27        df = pd.read_csv(csv_path)
 28        df[split] = df[split].apply(lambda x: json.loads(x.replace("'", '"')))  # ensures all items from column in list.
 29        split_list = df.iloc[0][split]
 30
 31    else:
 32        print(f"Creating a new split file at '{csv_path}'.")
 33        image_names = [
 34            os.path.basename(image).split(".")[0] for image in glob(os.path.join(path, data_dir, '*.tif'))
 35        ]
 36
 37        # Create random splits per dataset.
 38        train_ids, test_ids = train_test_split(image_names, test_size=0.2)  # 20% for test split.
 39        train_ids, val_ids = train_test_split(train_ids, test_size=0.15)  # 15% for val split.
 40        split_ids = {"train": train_ids, "val": val_ids, "test": test_ids}
 41
 42        df = pd.DataFrame.from_dict([split_ids])
 43        df.to_csv(csv_path, index=False)
 44
 45        split_list = split_ids[split]
 46
 47    return split_list
 48
 49
 50def get_cryonuseg_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 51    """Download the CryoNuSeg dataset for nucleus segmentation.
 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 folder where the data is downloaded and preprocessed.
 59    """
 60    data_dir = os.path.join(path, r"tissue images")
 61    if os.path.exists(os.path.join(path, r"tissue images")):
 62        return data_dir
 63
 64    os.makedirs(path, exist_ok=True)
 65    util.download_source_kaggle(
 66        path=path, dataset_name="ipateam/segmentation-of-nuclei-in-cryosectioned-he-images", download=download
 67    )
 68
 69    zip_path = os.path.join(path, "segmentation-of-nuclei-in-cryosectioned-he-images.zip")
 70    util.unzip(zip_path=zip_path, dst=path)
 71
 72    return data_dir
 73
 74
 75def get_cryonuseg_paths(
 76    path: Union[os.PathLike, str],
 77    split: Literal["train", "val", "test"],
 78    rater_choice: Literal["b1", "b2", "b3"] = "b1",
 79    download: bool = False,
 80) -> Tuple[List[str], List[str]]:
 81    """Get paths to the CryoNuSeg data.
 82
 83    Args:
 84        path: Filepath to a folder where the downloaded data will be saved.
 85        split: The choice of data split.
 86        rater: The choice of annotator.
 87        download: Whether to download the data if it is not present.
 88
 89    Returns:
 90        List of filepaths to the image data.
 91        List of filepaths to the label data.
 92    """
 93    data_dir = get_cryonuseg_data(path, download)
 94
 95    if rater_choice == "b1":
 96        label_dir = r"Annotator 1 (biologist)/"
 97    elif rater_choice == "b2":
 98        label_dir = r"Annotator 1 (biologist second round of manual marks up)/" * 2
 99    elif rater_choice == "b3":
100        label_dir = r"Annotator 2 (bioinformatician)/" * 2
101    else:
102        raise ValueError(f"'{rater_choice}' is not a valid rater choice.")
103
104    # Point to the instance labels folder
105    label_dir += r"label masks modify"
106    split_list = _create_split_csv(path, label_dir, split)
107
108    # Get the raw and label paths
109    label_paths = natsorted([os.path.join(path, label_dir, f'{fname}.tif') for fname in split_list])
110    raw_paths = natsorted([os.path.join(data_dir, f'{fname}.tif') for fname in split_list])
111
112    assert len(raw_paths) == len(label_paths) and len(raw_paths) > 0
113
114    return raw_paths, label_paths
115
116
117def get_cryonuseg_dataset(
118    path: Union[os.PathLike, str],
119    patch_shape: Tuple[int, int],
120    split: Literal["train", "val", "test"],
121    rater: Literal["b1", "b2", "b3"] = "b1",
122    resize_inputs: bool = False,
123    download: bool = False,
124    **kwargs
125) -> Dataset:
126    """Get the CryoNuSeg dataset for nucleus segmentation.
127
128    Args:
129        path: Filepath to a folder where the downloaded data will be saved.
130        patch_shape: The patch shape to use for training.
131        split: The choice of data split.
132        rater: The choice of annotator.
133        resize_inputs: Whether to resize the inputs.
134        download: Whether to download the data if it is not present.
135        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
136
137    Returns:
138        The segmentation dataset.
139    """
140    raw_paths, label_paths = get_cryonuseg_paths(path, split, rater, download)
141
142    if resize_inputs:
143        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": True}
144        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
145            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
146        )
147
148    return torch_em.default_segmentation_dataset(
149        raw_paths=raw_paths,
150        raw_key=None,
151        label_paths=label_paths,
152        label_key=None,
153        is_seg_dataset=False,
154        patch_shape=patch_shape,
155        **kwargs
156    )
157
158
159def get_cryonuseg_loader(
160    path: Union[os.PathLike, str],
161    batch_size: int,
162    patch_shape: Tuple[int, int],
163    split: Literal["train", "val", "test"],
164    rater: Literal["b1", "b2", "b3"] = "b1",
165    resize_inputs: bool = False,
166    download: bool = False,
167    **kwargs
168) -> DataLoader:
169    """Get the CryoNuSeg dataloader for nucleus segmentation.
170
171    Args:
172        path: Filepath to a folder where the downloaded data will be saved.
173        batch_size: The batch size for training.
174        patch_shape: The patch shape to use for training.
175        split: The choice of data split.
176        rater: The choice of annotator.
177        resize_inputs: Whether to resize the inputs.
178        download: Whether to download the data if it is not present.
179        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
180
181    Returns:
182        The DataLoader.
183    """
184    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
185    dataset = get_cryonuseg_dataset(path, patch_shape, split, rater, resize_inputs, download, **ds_kwargs)
186    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)