torch_em.data.datasets.histopathology.liver_tme_mibi

  1"""This dataset contains cell instance segmentation annotations for multiplexed ion beam imaging
  2(MIBI) of syngeneic orthotopic murine liver cancer models, covering intrahepatic cholangiocarcinoma
  3(iCCA) and hepatocellular carcinoma (HCC) driven by Trp53del / KrasG12D mutations, as well as
  4FGFR2 fusion-driven iCCA.
  5
  6The data is from the study "Multidimensional spatial profiling of the tumor microenvironment in
  7syngeneic orthotopic murine liver cancer models", hosted on the BioImage Archive at
  8https://www.ebi.ac.uk/biostudies/bioimages/studies/S-BIAD2557. Please cite it if you use this
  9dataset in your research.
 10
 11This loader covers the 562 fields of view (FOVs). Each FOV is a 1024x1024 crop with 30 protein
 12marker channels (covering cancer, immune and stromal lineages) and a matching per-cell instance
 13segmentation mask, generated with Cellpose on a nuclear (histone H3) and membrane marker composite
 14and filtered to high-confidence single cells. NOTE: The full set of FOVs is very large
 15(the 30-channel images alone total about 140 GB); use the `fovs` argument to restrict which ones
 16are downloaded and preprocessed.
 17
 18On first use, each requested FOV is converted into a single HDF5 file with the following layout:
 19    - 'raw/all': the (30, H, W) stack of all channels.
 20    - 'raw/channels/<channel>': each individual channel (H, W), see `CHANNELS` for the full list.
 21    - 'labels/instances': the instance segmentation.
 22"""
 23
 24import os
 25from glob import glob
 26from typing import List, Optional, Sequence, Tuple, Union
 27
 28import numpy as np
 29
 30from torch.utils.data import Dataset, DataLoader
 31
 32import torch_em
 33
 34from .. import util
 35
 36
 37BASE_URL = "https://ftp.ebi.ac.uk/pub/databases/biostudies/S-BIAD/557/S-BIAD2557/Files/spatial_murine_iCCAvsHCC"
 38MANIFEST_URL = f"{BASE_URL}/MIBI_images_file_list.json"
 39
 40CHANNELS = (
 41    "ATP5A", "B220", "cCASP3", "CD11c", "CD163", "CD31", "CD3e", "CD45", "CD4", "COL1A1", "CTLA4",
 42    "CytC", "F4_80", "G6PD", "GLUT1", "HH3", "HNF4a", "Ki67", "LDH", "Ly6G", "mCD11b", "mCD44",
 43    "mCD8", "mFoxP3", "mPD_1", "mPD_L1", "NaKATPase", "PanCK", "SMA", "Vimentin",
 44)
 45
 46
 47def _get_manifest(path, download):
 48    import requests
 49
 50    manifest_path = os.path.join(path, "MIBI_images_file_list.json")
 51    if not os.path.exists(manifest_path):
 52        if not download:
 53            raise RuntimeError(f"Cannot find the manifest at {manifest_path}, but download was set to False")
 54        response = requests.get(MANIFEST_URL)
 55        response.raise_for_status()
 56        with open(manifest_path, "wb") as f:
 57            f.write(response.content)
 58
 59    import json
 60    with open(manifest_path) as f:
 61        manifest = json.load(f)
 62
 63    return sorted(os.path.basename(entry["path"]) for entry in manifest)
 64
 65
 66def _convert_fov(fov_id, output_path, download):
 67    import tifffile
 68
 69    tmp_path = output_path + ".tmp"
 70
 71    channel_stack = []
 72    for channel in CHANNELS:
 73        channel_url = f"{BASE_URL}/image_data/{fov_id}/{channel}.tiff"
 74        channel_path = tmp_path + f".{channel}.tiff"
 75        util.download_source(channel_path, channel_url, download, checksum=None)
 76        channel_stack.append(tifffile.imread(channel_path))
 77        os.remove(channel_path)
 78    raw = np.stack(channel_stack, axis=0)
 79
 80    mask_url = f"{BASE_URL}/segmentation/cleaned_mask/{fov_id}_cleaned_mask.tiff"
 81    mask_path = tmp_path + ".mask.tiff"
 82    util.download_source(mask_path, mask_url, download, checksum=None)
 83    instances = tifffile.imread(mask_path)
 84    os.remove(mask_path)
 85
 86    import h5py
 87    with h5py.File(tmp_path, "w") as f:
 88        f.create_dataset("raw/all", data=raw, compression="gzip", chunks=(len(CHANNELS), 512, 512))
 89        for i, name in enumerate(CHANNELS):
 90            f.create_dataset(f"raw/channels/{name}", data=raw[i], compression="gzip")
 91        f.create_dataset("labels/instances", data=instances, compression="gzip")
 92
 93    os.replace(tmp_path, output_path)
 94
 95
 96def get_liver_tme_mibi_data(
 97    path: Union[os.PathLike, str], fovs: Optional[Sequence[str]] = None, download: bool = False,
 98) -> str:
 99    """Download and preprocess the murine liver TME MIBI data.
100
101    Args:
102        path: Filepath to a folder where the downloaded data will be saved.
103        fovs: The field of view (FOV) ids to prepare. By default all 562 FOVs are prepared, which
104            requires downloading about 140 GB of data. See `MANIFEST_URL` for the full list of ids.
105        download: Whether to download the data if it is not present.
106
107    Returns:
108        Filepath to the folder where the preprocessed data is stored.
109    """
110    os.makedirs(path, exist_ok=True)
111    valid_fovs = _get_manifest(path, download)
112
113    if fovs is None:
114        fovs = valid_fovs
115    else:
116        invalid = sorted(set(fovs) - set(valid_fovs))
117        if invalid:
118            raise ValueError(f"Invalid FOV id(s) {invalid}.")
119
120    preprocessed_dir = os.path.join(path, "preprocessed")
121    os.makedirs(preprocessed_dir, exist_ok=True)
122
123    for fov_id in fovs:
124        output_path = os.path.join(preprocessed_dir, f"{fov_id}.h5")
125        if os.path.exists(output_path):
126            continue
127        _convert_fov(fov_id, output_path, download)
128
129    return preprocessed_dir
130
131
132def get_liver_tme_mibi_paths(
133    path: Union[os.PathLike, str], fovs: Optional[Sequence[str]] = None, download: bool = False,
134) -> List[str]:
135    """Get paths to the preprocessed murine liver TME MIBI data.
136
137    Args:
138        path: Filepath to a folder where the downloaded data will be saved.
139        fovs: The field of view (FOV) ids to load. By default all 562 FOVs are loaded.
140        download: Whether to download the data if it is not present.
141
142    Returns:
143        List of filepaths to the preprocessed HDF5 files.
144    """
145    preprocessed_dir = get_liver_tme_mibi_data(path, fovs, download)
146    if fovs is None:
147        paths = sorted(glob(os.path.join(preprocessed_dir, "*.h5")))
148    else:
149        paths = [os.path.join(preprocessed_dir, f"{fov_id}.h5") for fov_id in fovs]
150
151    missing = [p for p in paths if not os.path.exists(p)]
152    if missing:
153        raise RuntimeError(f"Could not find the data at {missing}.")
154
155    return paths
156
157
158def get_liver_tme_mibi_dataset(
159    path: Union[os.PathLike, str],
160    patch_shape: Tuple[int, int],
161    fovs: Optional[Sequence[str]] = None,
162    channel: str = "all",
163    download: bool = False,
164    resize_inputs: bool = False,
165    **kwargs
166) -> Dataset:
167    """Get the murine liver TME MIBI dataset for cell instance segmentation in multiplexed
168    ion beam images of the liver cancer tumor microenvironment.
169
170    Args:
171        path: Filepath to a folder where the downloaded data will be saved.
172        patch_shape: The patch shape to use for training.
173        fovs: The field of view (FOV) ids to load. By default all 562 FOVs are loaded.
174        channel: The raw input. Either 'all' for the full (30, H, W) channel stack, or the name of
175            a single channel, see `CHANNELS` for the full list, e.g. 'PanCK'.
176        download: Whether to download the data if it is not present.
177        resize_inputs: Whether to resize the input images.
178        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
179
180    Returns:
181        The segmentation dataset.
182    """
183    if channel == "all":
184        raw_key, with_channels = "raw/all", True
185    elif channel in CHANNELS:
186        raw_key, with_channels = f"raw/channels/{channel}", False
187    else:
188        raise ValueError(f"'{channel}' is not a valid channel. Choose 'all' or one of {CHANNELS}.")
189
190    paths = get_liver_tme_mibi_paths(path, fovs, download)
191
192    if resize_inputs:
193        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
194        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
195            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
196        )
197
198    return torch_em.default_segmentation_dataset(
199        raw_paths=paths,
200        raw_key=raw_key,
201        label_paths=paths,
202        label_key="labels/instances",
203        patch_shape=patch_shape,
204        is_seg_dataset=True,
205        with_channels=with_channels,
206        ndim=2,
207        **kwargs
208    )
209
210
211def get_liver_tme_mibi_loader(
212    path: Union[os.PathLike, str],
213    patch_shape: Tuple[int, int],
214    batch_size: int,
215    fovs: Optional[Sequence[str]] = None,
216    channel: str = "all",
217    download: bool = False,
218    resize_inputs: bool = False,
219    **kwargs
220) -> DataLoader:
221    """Get the murine liver TME MIBI dataloader for cell instance segmentation in multiplexed
222    ion beam images of the liver cancer tumor microenvironment.
223
224    Args:
225        path: Filepath to a folder where the downloaded data will be saved.
226        patch_shape: The patch shape to use for training.
227        batch_size: The batch size for training.
228        fovs: The field of view (FOV) ids to load. By default all 562 FOVs are loaded.
229        channel: The raw input. Either 'all' for the full (30, H, W) channel stack, or the name of
230            a single channel, see `CHANNELS` for the full list, e.g. 'PanCK'.
231        download: Whether to download the data if it is not present.
232        resize_inputs: Whether to resize the input images.
233        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the
234            PyTorch DataLoader.
235
236    Returns:
237        The DataLoader.
238    """
239    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
240    dataset = get_liver_tme_mibi_dataset(
241        path, patch_shape, fovs=fovs, channel=channel, download=download, resize_inputs=resize_inputs, **ds_kwargs
242    )
243    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)