torch_em.data.datasets.histopathology.lipid_droplet_histology

  1"""The Lipid Droplet Histology dataset contains annotations for lipid droplet (fat vacuole) segmentation
  2in H&E stained histological images of liver tissue.
  3
  4This loader uses the "Corrected Subset B" of the dataset: 700 image patches of size 256 x 256 with manually
  5corrected masks of lipid droplets. The masks of the full whole-slide images (Datasets A and B) are not used,
  6since they require the whole-slide images, which are hosted elsewhere.
  7
  8The dataset is located at https://zenodo.org/records/7802211.
  9This dataset is from the publication https://openreview.net/forum?id=nTnAm_El0RC.
 10Please cite it if you use this dataset for your research.
 11"""
 12
 13import os
 14from glob import glob
 15from typing import Tuple, Union, List
 16
 17import numpy as np
 18import imageio.v3 as imageio
 19
 20from torch.utils.data import Dataset, DataLoader
 21
 22import torch_em
 23
 24from .. import util
 25
 26
 27URL = "https://zenodo.org/api/records/7802211/files/Corrected_Subset_B.zip/content"
 28CHECKSUM = "dc64f674f644f897024a295e23da95b52d82eecb2723a7ee1a3439c9f519ec30"
 29
 30
 31def _convert_to_tif(data_dir, image_dir, label_dir):
 32    import nibabel as nib
 33
 34    os.makedirs(image_dir, exist_ok=True)
 35    os.makedirs(label_dir, exist_ok=True)
 36
 37    src_dir = os.path.join(data_dir, "Corrected_Subset_B")
 38    label_paths = sorted(glob(os.path.join(src_dir, "labelsTs_subsetC_B", "*.nii.gz")))
 39    for label_path in label_paths:
 40        name = os.path.basename(label_path)[:-len(".nii.gz")]
 41
 42        # The red, green and blue channels are stored in separate files (suffix 0000, 0001 and 0002).
 43        channels = [
 44            np.asarray(nib.load(os.path.join(src_dir, "imagesTs_subsetC_B", f"{name}_000{i}.nii.gz")).dataobj)
 45            for i in range(3)
 46        ]
 47        image = np.stack([ch[..., 0].T for ch in channels], axis=-1).astype("uint8")
 48        label = (np.asarray(nib.load(label_path).dataobj)[..., 0].T > 0).astype("uint8")
 49
 50        imageio.imwrite(os.path.join(image_dir, f"{name}.tif"), image, compression="zlib")
 51        imageio.imwrite(os.path.join(label_dir, f"{name}.tif"), label, compression="zlib")
 52
 53
 54def get_lipid_droplet_histology_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 55    """Download the Lipid Droplet Histology dataset.
 56
 57    Args:
 58        path: Filepath to a folder where the downloaded data will be saved.
 59        download: Whether to download the data if it is not present.
 60
 61    Returns:
 62        Filepath where the data is downloaded.
 63    """
 64    data_dir = os.path.join(path, "data")
 65    image_dir, label_dir = os.path.join(path, "images"), os.path.join(path, "labels")
 66    if os.path.exists(image_dir) and os.path.exists(label_dir):
 67        return path
 68
 69    os.makedirs(path, exist_ok=True)
 70
 71    zip_path = os.path.join(path, "Corrected_Subset_B.zip")
 72    util.download_source(path=zip_path, url=URL, download=download, checksum=CHECKSUM)
 73    util.unzip(zip_path=zip_path, dst=data_dir)
 74
 75    _convert_to_tif(data_dir, image_dir, label_dir)
 76
 77    return path
 78
 79
 80def get_lipid_droplet_histology_paths(
 81    path: Union[os.PathLike, str], download: bool = False
 82) -> Tuple[List[str], List[str]]:
 83    """Get paths to the Lipid Droplet Histology data.
 84
 85    Args:
 86        path: Filepath to a folder where the downloaded data will be saved.
 87        download: Whether to download the data if it is not present.
 88
 89    Returns:
 90        List of filepaths for the image data.
 91        List of filepaths for the label data.
 92    """
 93    get_lipid_droplet_histology_data(path, download)
 94
 95    raw_paths = sorted(glob(os.path.join(path, "images", "*.tif")))
 96    label_paths = sorted(glob(os.path.join(path, "labels", "*.tif")))
 97    assert len(raw_paths) > 0 and len(raw_paths) == len(label_paths)
 98
 99    return raw_paths, label_paths
100
101
102def get_lipid_droplet_histology_dataset(
103    path: Union[os.PathLike, str],
104    patch_shape: Tuple[int, int],
105    download: bool = False,
106    **kwargs
107) -> Dataset:
108    """Get the Lipid Droplet Histology dataset for lipid droplet segmentation.
109
110    Args:
111        path: Filepath to a folder where the downloaded data will be saved.
112        patch_shape: The patch shape to use for training.
113        download: Whether to download the data if it is not present.
114        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
115
116    Returns:
117        The segmentation dataset.
118    """
119    raw_paths, label_paths = get_lipid_droplet_histology_paths(path, download)
120
121    return torch_em.default_segmentation_dataset(
122        raw_paths=raw_paths,
123        raw_key=None,
124        label_paths=label_paths,
125        label_key=None,
126        is_seg_dataset=False,
127        patch_shape=patch_shape,
128        ndim=2,
129        with_channels=True,
130        **kwargs
131    )
132
133
134def get_lipid_droplet_histology_loader(
135    path: Union[os.PathLike, str],
136    batch_size: int,
137    patch_shape: Tuple[int, int],
138    download: bool = False,
139    **kwargs
140) -> DataLoader:
141    """Get the Lipid Droplet Histology dataloader for lipid droplet segmentation.
142
143    Args:
144        path: Filepath to a folder where the downloaded data will be saved.
145        batch_size: The batch size for training.
146        patch_shape: The patch shape to use for training.
147        download: Whether to download the data if it is not present.
148        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
149
150    Returns:
151        The DataLoader.
152    """
153    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
154    dataset = get_lipid_droplet_histology_dataset(path, patch_shape, download, **ds_kwargs)
155    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)