torch_em.data.datasets.histopathology.tcga_tissue

  1"""The TCGA Tissue Segmentation dataset contains annotations for tissue vs. background
  2segmentation in histopathology slide images drawn from The Cancer Genome Atlas (TCGA).
  3
  4Each image is a full TCGA slide (predominantly H&E stained, some FFPE and frozen sections)
  5downsampled to 10 micrometer per pixel. The masks are binary, with tissue regions labeled as
  6foreground, including artifact-affected regions such as pen markings, ink, air bubbles and cracks.
  7
  8The dataset is located at https://huggingface.co/datasets/conflux-xyz/tcga-tissue-segmentation
  9and is licensed under CC0-1.0. It is not associated with a peer-reviewed publication. Users are
 10asked to cite the TCGA Research Network: https://www.cancer.gov/tcga.
 11"""
 12
 13import os
 14from glob import glob
 15from natsort import natsorted
 16from typing import List, Literal, Tuple, Union
 17
 18from torch.utils.data import Dataset, DataLoader
 19
 20import torch_em
 21
 22from .. import util
 23
 24
 25HF_REPO = "conflux-xyz/tcga-tissue-segmentation"
 26
 27
 28def get_tcga_tissue_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 29    """Download the TCGA Tissue Segmentation dataset.
 30
 31    Args:
 32        path: Filepath to a folder where the downloaded data will be saved.
 33        download: Whether to download the data if it is not present.
 34
 35    Returns:
 36        The filepath to the folder where the data is stored.
 37    """
 38    if os.path.exists(os.path.join(path, "images")) and os.path.exists(os.path.join(path, "masks")):
 39        return path
 40
 41    if not download:
 42        raise RuntimeError(f"Cannot find the data at {path}, but 'download' is set to False.")
 43
 44    try:
 45        from huggingface_hub import snapshot_download
 46    except ImportError:
 47        raise ImportError("huggingface_hub is required. Install with: pip install huggingface_hub")
 48
 49    os.makedirs(path, exist_ok=True)
 50    snapshot_download(repo_id=HF_REPO, repo_type="dataset", local_dir=path)
 51
 52    return path
 53
 54
 55def get_tcga_tissue_paths(
 56    path: Union[os.PathLike, str], split: Literal["train", "test"], download: bool = False,
 57) -> Tuple[List[str], List[str]]:
 58    """Get paths to the TCGA Tissue Segmentation data.
 59
 60    Args:
 61        path: Filepath to a folder where the downloaded data will be saved.
 62        split: The choice of data split, either 'train' or 'test'.
 63        download: Whether to download the data if it is not present.
 64
 65    Returns:
 66        List of filepaths for the image data.
 67        List of filepaths for the label data.
 68    """
 69    data_dir = get_tcga_tissue_data(path, download)
 70
 71    split_path = os.path.join(data_dir, f"{split}-slides.txt")
 72    with open(split_path) as f:
 73        slide_ids = {line.strip() for line in f if line.strip()}
 74
 75    raw_paths = natsorted(
 76        p for p in glob(os.path.join(data_dir, "images", "*.png"))
 77        if os.path.splitext(os.path.basename(p))[0] in slide_ids
 78    )
 79    label_paths = natsorted(
 80        p for p in glob(os.path.join(data_dir, "masks", "*.png"))
 81        if os.path.splitext(os.path.basename(p))[0] in slide_ids
 82    )
 83
 84    assert len(raw_paths) == len(label_paths) == len(slide_ids)
 85    assert all(
 86        os.path.basename(raw_path) == os.path.basename(label_path)
 87        for raw_path, label_path in zip(raw_paths, label_paths)
 88    )
 89
 90    return raw_paths, label_paths
 91
 92
 93def get_tcga_tissue_dataset(
 94    path: Union[os.PathLike, str],
 95    patch_shape: Tuple[int, int],
 96    split: Literal["train", "test"] = "train",
 97    resize_inputs: bool = False,
 98    download: bool = False,
 99    **kwargs,
100) -> Dataset:
101    """Get the TCGA Tissue Segmentation dataset for tissue vs. background segmentation.
102
103    Args:
104        path: Filepath to a folder where the downloaded data will be saved.
105        patch_shape: The patch shape to use for training.
106        split: The choice of data split, either 'train' or 'test'.
107        resize_inputs: Whether to resize the inputs.
108        download: Whether to download the data if it is not present.
109        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
110
111    Returns:
112        The segmentation dataset.
113    """
114    raw_paths, label_paths = get_tcga_tissue_paths(path, split, download)
115
116    if resize_inputs:
117        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": True}
118        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
119            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
120        )
121
122    return torch_em.default_segmentation_dataset(
123        raw_paths=raw_paths,
124        raw_key=None,
125        label_paths=label_paths,
126        label_key=None,
127        patch_shape=patch_shape,
128        is_seg_dataset=False,
129        ndim=2,
130        with_channels=True,
131        **kwargs,
132    )
133
134
135def get_tcga_tissue_loader(
136    path: Union[os.PathLike, str],
137    batch_size: int,
138    patch_shape: Tuple[int, int],
139    split: Literal["train", "test"] = "train",
140    resize_inputs: bool = False,
141    download: bool = False,
142    **kwargs,
143) -> DataLoader:
144    """Get the TCGA Tissue Segmentation dataloader for tissue vs. background segmentation.
145
146    Args:
147        path: Filepath to a folder where the downloaded data will be saved.
148        batch_size: The batch size for training.
149        patch_shape: The patch shape to use for training.
150        split: The choice of data split, either 'train' or 'test'.
151        resize_inputs: Whether to resize the inputs.
152        download: Whether to download the data if it is not present.
153        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or the PyTorch DataLoader.
154
155    Returns:
156        The DataLoader.
157    """
158    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
159    dataset = get_tcga_tissue_dataset(path, patch_shape, split, resize_inputs, download, **ds_kwargs)
160    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)