torch_em.data.datasets.light_microscopy.hipsc_single_cell

  1"""The hiPSC single-cell image dataset contains 3D confocal fluorescence microscopy images of
  2human induced pluripotent stem cells (hiPSC), with segmentation masks for the nucleus, cell
  3membrane, and one fluorescently-tagged subcellular structure per cell line (e.g. mitochondria,
  4Golgi, microtubules).
  5
  6The dataset provides per-cell crops (binary masks) and full field-of-view images (instance masks
  7for nucleus and cell, binary mask for the structure). It is located at
  8https://open.quiltdata.com/b/allencell/packages/aics/hipsc_single_cell_image_dataset under the
  9Allen Institute for Cell Science Terms of Use (https://www.allencell.org/terms-of-use.html).
 10This dataset is from the publication https://doi.org/10.1038/s41586-022-05563-7.
 11Please cite it if you use this dataset in your research.
 12
 13NOTE: The full dataset covers almost 32,000 cells and is roughly a terabyte in size. Use
 14`structure_names` and/or `n_samples` to bound the download to a manageable subset.
 15"""
 16
 17import os
 18from typing import List, Literal, Optional, Tuple, Union
 19
 20import numpy as np
 21
 22from torch.utils.data import Dataset, DataLoader
 23
 24import torch_em
 25
 26from .. import util
 27
 28
 29BASE_URL = "https://allencell.s3.amazonaws.com/aics/hipsc_single_cell_image_dataset"
 30
 31# Crop-level segmentation channels, as documented in the package README's `name_dict`:
 32# ['dna_segmentation', 'membrane_segmentation', 'membrane_segmentation_roof',
 33#  'struct_segmentation', 'struct_segmentation_roof'].
 34CROP_SEG_CHANNELS = {"nucleus": 0, "cell": 1, "structure": 3}
 35# FOV-level segmentation channels, as documented in the package README's `fov_seg_path`:
 36# nuclear segmentation, cell segmentation, contour of nuclei, contour of cell.
 37FOV_SEG_CHANNELS = {"nucleus": 0, "cell": 1}
 38
 39VALID_TARGETS = ["nucleus", "cell", "structure"]
 40VALID_SAMPLE_TYPES = ["cell", "fov"]
 41
 42
 43def _get_metadata(path, download):
 44    import pandas as pd
 45
 46    csv_path = os.path.join(path, "metadata.csv")
 47    util.download_source(path=csv_path, url=f"{BASE_URL}/metadata.csv", download=download, checksum=None)
 48    return pd.read_csv(csv_path)
 49
 50
 51def _select_rows(path, structure_names, sample_type, n_samples, download):
 52    df = _get_metadata(path, download)
 53    if structure_names is not None:
 54        df = df[df["structure_name"].isin(structure_names)]
 55    if n_samples is not None:
 56        df = df.groupby("structure_name", group_keys=False).head(n_samples)
 57    if sample_type == "fov":
 58        df = df.groupby("FOVId", as_index=False).first()
 59    return df.reset_index(drop=True)
 60
 61
 62def _download_cell_files(path, row, download):
 63    raw_path = os.path.join(path, row["crop_raw"])
 64    seg_path = os.path.join(path, row["crop_seg"])
 65    os.makedirs(os.path.dirname(raw_path), exist_ok=True)
 66    os.makedirs(os.path.dirname(seg_path), exist_ok=True)
 67    util.download_source(path=raw_path, url=f"{BASE_URL}/{row['crop_raw']}", download=download, checksum=None)
 68    util.download_source(path=seg_path, url=f"{BASE_URL}/{row['crop_seg']}", download=download, checksum=None)
 69
 70
 71def _download_fov_files(path, row, download):
 72    raw_path = os.path.join(path, row["fov_path"])
 73    seg_path = os.path.join(path, row["fov_seg_path"])
 74    struct_path = os.path.join(path, row["struct_seg_path"])
 75    for out_path in (raw_path, seg_path, struct_path):
 76        os.makedirs(os.path.dirname(out_path), exist_ok=True)
 77    util.download_source(path=raw_path, url=f"{BASE_URL}/{row['fov_path']}", download=download, checksum=None)
 78    util.download_source(path=seg_path, url=f"{BASE_URL}/{row['fov_seg_path']}", download=download, checksum=None)
 79    util.download_source(
 80        path=struct_path, url=f"{BASE_URL}/{row['struct_seg_path']}", download=download, checksum=None
 81    )
 82
 83
 84def _create_h5(path, df, sample_type, target):
 85    import h5py
 86    import tifffile
 87    from tqdm import tqdm
 88
 89    h5_dir = os.path.join(path, "h5_data", sample_type, target)
 90    os.makedirs(h5_dir, exist_ok=True)
 91
 92    h5_paths = []
 93    for _, row in tqdm(df.iterrows(), total=len(df), desc=f"Preparing '{sample_type}/{target}' data"):
 94        h5_path = os.path.join(h5_dir, f"{row['CellId']}.h5")
 95        h5_paths.append(h5_path)
 96        if os.path.exists(h5_path):
 97            continue
 98
 99        if sample_type == "cell":
100            raw = tifffile.imread(os.path.join(path, row["crop_raw"]))  # (Z, 3, Y, X): dna, membrane, structure
101            seg = tifffile.imread(os.path.join(path, row["crop_seg"]))  # (Z, 5, Y, X), see CROP_SEG_CHANNELS
102            raw = np.moveaxis(raw, 1, 0)
103            labels = (seg[:, CROP_SEG_CHANNELS[target]] > 0).astype("uint8")
104        else:
105            raw_full = tifffile.imread(os.path.join(path, row["fov_path"]))  # (Z, 7, Y, X)
106            channel_idx = {
107                "dna": row["ChannelNumber405"],
108                "membrane": row["ChannelNumber638"],
109                "structure": row["ChannelNumberStruct"],
110            }
111            raw = np.stack([raw_full[:, channel_idx[name]] for name in ("dna", "membrane", "structure")], axis=0)
112            if target == "structure":
113                labels = (tifffile.imread(os.path.join(path, row["struct_seg_path"])) > 0).astype("uint8")
114            else:
115                seg = tifffile.imread(os.path.join(path, row["fov_seg_path"]))  # (Z, 4, Y, X), see FOV_SEG_CHANNELS
116                labels = seg[:, FOV_SEG_CHANNELS[target]].astype("uint32")
117
118        with h5py.File(h5_path, "w") as f:
119            f.create_dataset("raw", data=raw, compression="gzip")
120            f.create_dataset("labels", data=labels, compression="gzip")
121
122    return h5_paths
123
124
125def get_hipsc_single_cell_data(
126    path: Union[os.PathLike, str],
127    structure_names: Optional[List[str]] = None,
128    sample_type: Literal["cell", "fov"] = "cell",
129    n_samples: Optional[int] = None,
130    download: bool = False,
131) -> str:
132    """Download the hiPSC single-cell image dataset.
133
134    Args:
135        path: Filepath to a folder where the downloaded data will be saved.
136        structure_names: The tagged structures to restrict the data to, e.g. ['TOMM20']. By default all are used.
137        sample_type: Whether to download per-cell crops or full field-of-view (FOV) images.
138        n_samples: The maximum number of cells to use per structure. By default all are used.
139        download: Whether to download the data if it is not present.
140
141    Returns:
142        The filepath to the folder with the downloaded data.
143    """
144    assert sample_type in VALID_SAMPLE_TYPES, f"'{sample_type}' is not a valid sample type: {VALID_SAMPLE_TYPES}."
145    os.makedirs(path, exist_ok=True)
146    df = _select_rows(path, structure_names, sample_type, n_samples, download)
147    for _, row in df.iterrows():
148        if sample_type == "cell":
149            _download_cell_files(path, row, download)
150        else:
151            _download_fov_files(path, row, download)
152    return path
153
154
155def get_hipsc_single_cell_paths(
156    path: Union[os.PathLike, str],
157    structure_names: Optional[List[str]] = None,
158    sample_type: Literal["cell", "fov"] = "cell",
159    target: Literal["nucleus", "cell", "structure"] = "nucleus",
160    n_samples: Optional[int] = None,
161    download: bool = False,
162) -> List[str]:
163    """Get paths to the hiPSC single-cell image data.
164
165    Args:
166        path: Filepath to a folder where the downloaded data will be saved.
167        structure_names: The tagged structures to restrict the data to, e.g. ['TOMM20']. By default all are used.
168        sample_type: Whether to use per-cell crops or full field-of-view (FOV) images.
169        target: The segmentation target. One of 'nucleus', 'cell' or 'structure'. For `sample_type='cell'`, all
170            targets are binary masks. For `sample_type='fov'`, 'nucleus' and 'cell' are instance masks and
171            'structure' is a binary mask.
172        n_samples: The maximum number of cells to use per structure. By default all are used.
173        download: Whether to download the data if it is not present.
174
175    Returns:
176        List of filepaths to the h5 files with the raw and label data.
177    """
178    assert target in VALID_TARGETS, f"'{target}' is not a valid target. Choose from {VALID_TARGETS}."
179
180    get_hipsc_single_cell_data(path, structure_names, sample_type, n_samples, download)
181    df = _select_rows(path, structure_names, sample_type, n_samples, download=False)
182    h5_paths = _create_h5(path, df, sample_type, target)
183
184    return h5_paths
185
186
187def get_hipsc_single_cell_dataset(
188    path: Union[os.PathLike, str],
189    patch_shape: Tuple[int, int, int],
190    structure_names: Optional[List[str]] = None,
191    sample_type: Literal["cell", "fov"] = "cell",
192    target: Literal["nucleus", "cell", "structure"] = "nucleus",
193    n_samples: Optional[int] = None,
194    download: bool = False,
195    **kwargs,
196) -> Dataset:
197    """Get the hiPSC single-cell image dataset for 3D segmentation.
198
199    Args:
200        path: Filepath to a folder where the downloaded data will be saved.
201        patch_shape: The patch shape to use for training.
202        structure_names: The tagged structures to restrict the data to, e.g. ['TOMM20']. By default all are used.
203        sample_type: Whether to use per-cell crops or full field-of-view (FOV) images.
204        target: The segmentation target. One of 'nucleus', 'cell' or 'structure'.
205        n_samples: The maximum number of cells to use per structure. By default all are used.
206        download: Whether to download the data if it is not present.
207        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
208
209    Returns:
210        The segmentation dataset.
211    """
212    h5_paths = get_hipsc_single_cell_paths(path, structure_names, sample_type, target, n_samples, download)
213
214    kwargs, _ = util.add_instance_label_transform(kwargs, add_binary_target=True)
215    kwargs = util.ensure_transforms(ndim=3, **kwargs)
216
217    return torch_em.default_segmentation_dataset(
218        raw_paths=h5_paths,
219        raw_key="raw",
220        label_paths=h5_paths,
221        label_key="labels",
222        patch_shape=patch_shape,
223        with_channels=True,
224        ndim=3,
225        **kwargs,
226    )
227
228
229def get_hipsc_single_cell_loader(
230    path: Union[os.PathLike, str],
231    batch_size: int,
232    patch_shape: Tuple[int, int, int],
233    structure_names: Optional[List[str]] = None,
234    sample_type: Literal["cell", "fov"] = "cell",
235    target: Literal["nucleus", "cell", "structure"] = "nucleus",
236    n_samples: Optional[int] = None,
237    download: bool = False,
238    **kwargs,
239) -> DataLoader:
240    """Get the hiPSC single-cell image dataloader for 3D segmentation.
241
242    Args:
243        path: Filepath to a folder where the downloaded data will be saved.
244        batch_size: The batch size for training.
245        patch_shape: The patch shape to use for training.
246        structure_names: The tagged structures to restrict the data to, e.g. ['TOMM20']. By default all are used.
247        sample_type: Whether to use per-cell crops or full field-of-view (FOV) images.
248        target: The segmentation target. One of 'nucleus', 'cell' or 'structure'.
249        n_samples: The maximum number of cells to use per structure. By default all are used.
250        download: Whether to download the data if it is not present.
251        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
252
253    Returns:
254        The DataLoader.
255    """
256    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
257    dataset = get_hipsc_single_cell_dataset(
258        path, patch_shape, structure_names, sample_type, target, n_samples, download, **ds_kwargs,
259    )
260    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)