torch_em.data.datasets.histopathology.bcss

  1"""This dataset contains annotations for tissue region segmentation in
  2breast cancer histopathology images.
  3
  4NOTE: There are multiple semantic instances in tissue labels. Below mentioned are their respective index details:
  5    - 0: outside_roi (~background)
  6    - 1: tumor
  7    - 2: stroma
  8    - 3: lymphocytic_infiltrate
  9    - 4: necrosis_or_debris
 10    - 5: glandular_secretions
 11    - 6: blood
 12    - 7: exclude
 13    - 8: metaplasia_NOS
 14    - 9: fat
 15    - 10: plasma_cells
 16    - 11: other_immune_infiltrate
 17    - 12: mucoid_material
 18    - 13: normal_acinus_or_duct
 19    - 14: lymphatics
 20    - 15: undetermined
 21    - 16: nerve
 22    - 17: skin_adnexa
 23    - 18: blood_vessel
 24    - 19: angioinvasion
 25    - 20: dcis
 26    - 21: other
 27
 28This dataset is from https://bcsegmentation.grand-challenge.org/BCSS/.
 29Please cite this paper (https://doi.org/10.1093/bioinformatics/btz083) if you use this dataset for a publication.
 30"""
 31
 32import os
 33import shutil
 34from glob import glob
 35from pathlib import Path
 36from warnings import warn
 37from typing import Union, Optional, List, Tuple
 38
 39from packaging import version
 40from sklearn.model_selection import train_test_split
 41
 42import torch
 43from torch.utils.data import Dataset, DataLoader
 44
 45import torch_em
 46
 47from .. import util
 48
 49try:
 50    import gdown
 51except ImportError:
 52    gdown = None
 53
 54
 55URL = "https://drive.google.com/drive/folders/1zqbdkQF8i5cEmZOGmbdQm-EP8dRYtvss?usp=sharing"
 56
 57
 58# TODO
 59CHECKSUM = None
 60
 61# Google Drive can deny access to some files in the folder above. A retry usually gets them.
 62MAX_DOWNLOAD_ATTEMPTS = 3
 63
 64
 65TEST_LIST = [
 66    "TCGA-A2-A0SX-DX1_xmin53791_ymin56683_MPP-0.2500", "TCGA-BH-A0BG-DX1_xmin64019_ymin24975_MPP-0.2500",
 67    "TCGA-AR-A1AI-DX1_xmin38671_ymin10616_MPP-0.2500", "TCGA-E2-A574-DX1_xmin54962_ymin47475_MPP-0.2500",
 68    "TCGA-GM-A3XL-DX1_xmin29910_ymin15820_MPP-0.2500", "TCGA-E2-A14X-DX1_xmin88836_ymin66393_MPP-0.2500",
 69    "TCGA-A2-A04P-DX1_xmin104246_ymin48517_MPP-0.2500", "TCGA-E2-A14N-DX1_xmin21383_ymin66838_MPP-0.2500",
 70    "TCGA-EW-A1OV-DX1_xmin126026_ymin65132_MPP-0.2500", "TCGA-S3-AA15-DX1_xmin55486_ymin28926_MPP-0.2500",
 71    "TCGA-LL-A5YO-DX1_xmin36631_ymin44396_MPP-0.2500", "TCGA-GI-A2C9-DX1_xmin20882_ymin11843_MPP-0.2500",
 72    "TCGA-BH-A0BW-DX1_xmin42346_ymin30843_MPP-0.2500", "TCGA-E2-A1B6-DX1_xmin16266_ymin50634_MPP-0.2500",
 73    "TCGA-AO-A0J2-DX1_xmin33561_ymin14515_MPP-0.2500"
 74]
 75
 76
 77def _download_bcss_dataset(path, download):
 78    for attempt in range(1, MAX_DOWNLOAD_ATTEMPTS + 1):
 79        if not os.path.exists(path):
 80            util.download_source_gdrive(
 81                path=path, url=URL, download=download, checksum=CHECKSUM, download_type="folder"
 82            )
 83        else:
 84            # `download_folder` skips files that already exist on disk.
 85            # This call only retries the files that a previous attempt could not download.
 86            assert version.parse(gdown.__version__) == version.parse("4.6.3"), "Please install 'gdown==4.6.3'."
 87            gdown.download_folder.__globals__["MAX_NUMBER_FILES"] = 10000
 88            gdown.download_folder(url=URL, output=path, quiet=True, remaining_ok=True)
 89
 90        n_images = len(glob(os.path.join(path, "rgbs_colorNormalized", "*")))
 91        n_masks = len(glob(os.path.join(path, "masks", "*")))
 92        if n_images > 0 and n_images == n_masks:
 93            return
 94        print(f"Download attempt {attempt} of {MAX_DOWNLOAD_ATTEMPTS} found {n_images} images and {n_masks} masks.")
 95
 96    print(
 97        "Google Drive did not serve every file after several attempts. "
 98        "The dataset will use only the images that have a matching mask."
 99    )
100
101
102def _get_image_and_label_paths(path):
103    # when downloading the files from `URL`, the input images are stored under `rgbs_colorNormalized`
104    # when getting the files from the git repo's command line feature, the input images are stored under `images`
105    if os.path.exists(os.path.join(path, "images")):
106        image_dir = os.path.join(path, "images")
107    elif os.path.exists(os.path.join(path, "rgbs_colorNormalized")):
108        image_dir = os.path.join(path, "rgbs_colorNormalized")
109    else:
110        raise ValueError(
111            "Please check the image directory. "
112            "If downloaded from gdrive, it's named \"rgbs_colorNormalized\", if from github it's named \"images\""
113        )
114    label_dir = os.path.join(path, "masks")
115
116    # Google Drive can deny access to individual files, so the folder download can skip some of them.
117    # We pair each image with its mask by filename, not by their sorted order.
118    image_stems = {Path(p).stem: p for p in glob(os.path.join(image_dir, "*"))}
119    label_stems = {Path(p).stem: p for p in glob(os.path.join(label_dir, "*"))}
120    common_stems = sorted(set(image_stems) & set(label_stems))
121
122    if len(common_stems) < len(image_stems) or len(common_stems) < len(label_stems):
123        warn(
124            f"Found {len(image_stems)} images and {len(label_stems)} masks, but only {len(common_stems)} "
125            "of them form a matching pair. The dataset will use only the matching pairs."
126        )
127
128    image_paths = [image_stems[stem] for stem in common_stems]
129    label_paths = [label_stems[stem] for stem in common_stems]
130
131    return image_paths, label_paths
132
133
134def get_bcss_data(path: Union[os.PathLike, str], download: bool = False):
135    """Download the BCSS dataset.
136
137    Args:
138        path: Filepath to a folder where the downloaded data will be saved.
139        download: Whether to download the data if it is not present.
140    """
141    if download:
142        _download_bcss_dataset(path, download)
143
144    if os.path.exists(os.path.join(path, "train")) and os.path.exists(os.path.join(path, "test")):
145        return
146
147    all_image_paths, all_label_paths = _get_image_and_label_paths(path)
148
149    train_img_dir, train_lab_dir = os.path.join(path, "train", "images"), os.path.join(path, "train", "masks")
150    test_img_dir, test_lab_dir = os.path.join(path, "test", "images"), os.path.join(path, "test", "masks")
151    os.makedirs(train_img_dir, exist_ok=True)
152    os.makedirs(train_lab_dir, exist_ok=True)
153    os.makedirs(test_img_dir, exist_ok=True)
154    os.makedirs(test_lab_dir, exist_ok=True)
155
156    for image_path, label_path in zip(all_image_paths, all_label_paths):
157        img_idx, label_idx = os.path.split(image_path)[-1], os.path.split(label_path)[-1]
158        if Path(image_path).stem in TEST_LIST:
159            # move image and label to test
160            dst_img_path, dst_lab_path = os.path.join(test_img_dir, img_idx), os.path.join(test_lab_dir, label_idx)
161            shutil.copy(src=image_path, dst=dst_img_path)
162            shutil.copy(src=label_path, dst=dst_lab_path)
163        else:
164            # move image and label to train
165            dst_img_path, dst_lab_path = os.path.join(train_img_dir, img_idx), os.path.join(train_lab_dir, label_idx)
166            shutil.copy(src=image_path, dst=dst_img_path)
167            shutil.copy(src=label_path, dst=dst_lab_path)
168
169
170def get_bcsss_paths(
171    path: Union[os.PathLike, str], split: Optional[str] = None, val_fraction: float = 0.2, download: bool = False
172) -> Tuple[List[str], List[str]]:
173    """Get paths to the BCSS data.
174
175    Args:
176        path: Filepath to a folder where the downloaded data will be saved.
177        split: The split to use for the dataset. Either 'train', 'val' or 'test'.
178        val_fraction: The fraction of data to be considered for validation split.
179        download: Whether to download the data if it is not present.
180
181    Returns:
182        List of filepaths for the image data.
183        List of filepaths for the label data.
184    """
185    get_bcss_data(path, download)
186
187    if split is None:
188        image_paths = sorted(glob(os.path.join(path, "*", "images", "*")))
189        label_paths = sorted(glob(os.path.join(path, "*", "masks", "*")))
190    else:
191        assert split in ["train", "val", "test"], "Please choose from the available `train` / `val` / `test` splits"
192        if split == "test":
193            image_paths = sorted(glob(os.path.join(path, "test", "images", "*")))
194            label_paths = sorted(glob(os.path.join(path, "test", "masks", "*")))
195        else:
196            image_paths = sorted(glob(os.path.join(path, "train", "images", "*")))
197            label_paths = sorted(glob(os.path.join(path, "train", "masks", "*")))
198
199            (train_image_paths, val_image_paths,
200             train_label_paths, val_label_paths) = train_test_split(
201                image_paths, label_paths, test_size=val_fraction, random_state=42
202            )
203
204            image_paths = train_image_paths if split == "train" else val_image_paths
205            label_paths = train_label_paths if split == "train" else val_label_paths
206
207    assert len(image_paths) == len(label_paths)
208
209    return image_paths, label_paths
210
211
212def get_bcss_dataset(
213    path: Union[os.PathLike, str],
214    patch_shape: Tuple[int, ...],
215    split: Optional[str] = None,
216    val_fraction: float = 0.2,
217    download: bool = False,
218    label_dtype: torch.dtype = torch.int64,
219    **kwargs
220) -> Dataset:
221    """Get the BCSS dataset for breast cancer tissue segmentation in histopathology.
222
223    Args:
224        path: Filepath to a folder where the downloaded data will be saved.
225        patch_shape: The patch shape to use for training.
226        split: The split to use for the dataset. Either 'train', 'val' or 'test'.
227        val_fraction: The fraction of data to be considered for validation split.
228        download: Whether to download the data if it is not present.
229        label_dtype: The datatype of labels.
230        kwargs: kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
231
232    Returns:
233        The segmentation dataset.
234    """
235    image_paths, label_paths = get_bcsss_paths(path, split, val_fraction, download)
236
237    return torch_em.default_segmentation_dataset(
238        raw_paths=image_paths,
239        raw_key=None,
240        label_paths=label_paths,
241        label_key=None,
242        patch_shape=patch_shape,
243        label_dtype=label_dtype,
244        is_seg_dataset=False,
245        **kwargs
246    )
247
248
249def get_bcss_loader(
250    path: Union[os.PathLike, str],
251    patch_shape: Tuple[int, ...],
252    batch_size: int,
253    split: Optional[str] = None,
254    val_fraction: float = 0.2,
255    download: bool = False,
256    label_dtype: torch.dtype = torch.int64,
257    **kwargs
258) -> DataLoader:
259    """Get the BCSS dataloader for breast cancer tissue segmentation in histopathology.
260
261    Args:
262        path: Filepath to a folder where the downloaded data will be saved.
263        patch_shape: The patch shape to use for training.
264        batch_size: The batch size for training.
265        split: The split to use for the dataset. Either 'train', 'val' or 'test'.
266        val_fraction: The fraction of data to be considered for validation split.
267        download: Whether to download the data if it is not present.
268        label_dtype: The datatype of labels.
269        kwargs: kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
270
271    Returns:
272        The DataLoader.
273    """
274    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
275    dataset = get_bcss_dataset(
276        path, patch_shape, split=split, val_fraction=val_fraction, download=download,
277        label_dtype=label_dtype, **ds_kwargs
278    )
279    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)