torch_em.data.datasets.light_microscopy.myofuse

  1"""MyoFuse contains annotations for myonuclei in fluorescence microscopy images of muscle cell cultures.
  2
  3The images show the myotube (myosin) channel of mouse C2C12 and human primary myotube cultures.
  4Every nucleus has a class. The class states if the nucleus is inside or outside a myotube.
  5MyoFuse uses these classes to measure the fusion index.
  6NOTE: A custom cellpose model created the nuclei instance masks. Nobody curated them by hand.
  7NOTE: This dataset does not contain the nuclei (DAPI) channel. Only the myotube channel has annotations.
  8
  9The dataset is located at https://doi.org/10.5281/zenodo.14731491.
 10This dataset is from the publication https://github.com/BenLair/MyoFuse.
 11Please cite it if you use this dataset for your research.
 12"""
 13
 14import os
 15from glob import glob
 16from natsort import natsorted
 17from typing import Union, Tuple, List, Literal, Optional
 18
 19import numpy as np
 20import tifffile
 21
 22from torch.utils.data import Dataset, DataLoader
 23
 24import torch_em
 25
 26from .. import util
 27
 28
 29URL = "https://zenodo.org/records/14731491/files/Training%20Images%20MyoFuse%20v1.0.0.zip"
 30CHECKSUM = "247725509cf5dd785ba919438d2eba7f544dc588880396f245b0700664bdd0b8"
 31
 32SUBSET_PREFIX = {"human": "H", "mouse": "M"}
 33
 34
 35def _get_semantic_labels(instances, regions, classes):
 36    """Map the class of every nucleus onto the instance mask.
 37
 38    MyoFuse stores the regions in a shuffled order. This function matches each region
 39    to its instance id with the coordinates of the region.
 40    """
 41    semantic = np.zeros(instances.shape, dtype="uint8")
 42    for region, class_id in zip(regions, classes):
 43        coords = region["coords"]
 44        instance_ids = np.unique(instances[coords[:, 0], coords[:, 1]])
 45        assert len(instance_ids) == 1, f"The region maps to more than one instance: {instance_ids}."
 46        semantic[instances == instance_ids[0]] = class_id + 1
 47
 48    return semantic
 49
 50
 51def _preprocess_data(input_dir, data_dir):
 52    import h5py
 53    import torch
 54
 55    os.makedirs(data_dir, exist_ok=True)
 56
 57    annotations = torch.load(os.path.join(input_dir, "labels"), weights_only=False)
 58    for i, image_path in enumerate(annotations["image_path"]):
 59        fname = os.path.basename(image_path)
 60        image = tifffile.imread(os.path.join(input_dir, "Images", fname))
 61        instances = tifffile.imread(os.path.join(input_dir, "Masks", fname)).astype("uint16")
 62        semantic = _get_semantic_labels(instances, annotations["regionprops"][i], annotations["labels_list"][i])
 63
 64        with h5py.File(os.path.join(data_dir, os.path.splitext(fname)[0] + ".h5"), "a") as f:
 65            f.create_dataset("raw", data=image, compression="gzip")
 66            f.create_dataset("labels/instances", data=instances, compression="gzip")
 67            f.create_dataset("labels/semantic", data=semantic, compression="gzip")
 68
 69
 70def get_myofuse_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 71    """Download the MyoFuse dataset.
 72
 73    Args:
 74        path: The folder where the function stores the data.
 75        download: Whether to download the data if it is not present.
 76
 77    Returns:
 78        The filepath to the folder with the prepared data.
 79    """
 80    data_dir = os.path.join(path, "data")
 81    if os.path.exists(data_dir):
 82        return data_dir
 83
 84    os.makedirs(path, exist_ok=True)
 85
 86    zip_path = os.path.join(path, "Training_Images_MyoFuse.zip")
 87    util.download_source(path=zip_path, url=URL, download=download, checksum=CHECKSUM)
 88    util.unzip(zip_path=zip_path, dst=path)
 89
 90    _preprocess_data(os.path.join(path, "Training Images MyoFuse v1.0.0"), data_dir)
 91
 92    return data_dir
 93
 94
 95def get_myofuse_paths(
 96    path: Union[os.PathLike, str],
 97    subset: Optional[Literal["human", "mouse"]] = None,
 98    download: bool = False,
 99) -> List[str]:
100    """Get the paths to the MyoFuse data.
101
102    Args:
103        path: The folder where the function stores the data.
104        subset: The cell type. Use 'human' for human primary myotubes or 'mouse' for mouse C2C12 myotubes.
105            The function uses both cell types by default.
106        download: Whether to download the data if it is not present.
107
108    Returns:
109        The list of filepaths to the input data.
110    """
111    data_dir = get_myofuse_data(path, download)
112
113    if subset is None:
114        pattern = "*.h5"
115    else:
116        if subset not in SUBSET_PREFIX:
117            raise ValueError(f"'{subset}' is not a valid subset. Choose one of {list(SUBSET_PREFIX.keys())}.")
118        pattern = f"{SUBSET_PREFIX[subset]}_*.h5"
119
120    volume_paths = natsorted(glob(os.path.join(data_dir, pattern)))
121    assert len(volume_paths) > 0, f"Could not find data for the subset '{subset}'."
122    return volume_paths
123
124
125def get_myofuse_dataset(
126    path: Union[os.PathLike, str],
127    patch_shape: Tuple[int, int],
128    subset: Optional[Literal["human", "mouse"]] = None,
129    label_choice: Literal["semantic", "instances"] = "semantic",
130    download: bool = False,
131    **kwargs
132) -> Dataset:
133    """Get the MyoFuse dataset for segmentation of myonuclei.
134
135    Args:
136        path: The folder where the function stores the data.
137        patch_shape: The patch shape to use for training.
138        subset: The cell type. Use 'human' for human primary myotubes or 'mouse' for mouse C2C12 myotubes.
139            The function uses both cell types by default.
140        label_choice: The label type. Use 'semantic' for the class of every nucleus,
141            or 'instances' for the masks of the nuclei.
142        download: Whether to download the data if it is not present.
143        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
144
145    Returns:
146        The segmentation dataset.
147    """
148    if label_choice not in ("semantic", "instances"):
149        raise ValueError(f"'{label_choice}' is not a valid label choice. Choose 'semantic' or 'instances'.")
150
151    volume_paths = get_myofuse_paths(path, subset, download)
152
153    return torch_em.default_segmentation_dataset(
154        raw_paths=volume_paths,
155        raw_key="raw",
156        label_paths=volume_paths,
157        label_key=f"labels/{label_choice}",
158        patch_shape=patch_shape,
159        is_seg_dataset=True,
160        ndim=2,
161        **kwargs
162    )
163
164
165def get_myofuse_loader(
166    path: Union[os.PathLike, str],
167    batch_size: int,
168    patch_shape: Tuple[int, int],
169    subset: Optional[Literal["human", "mouse"]] = None,
170    label_choice: Literal["semantic", "instances"] = "semantic",
171    download: bool = False,
172    **kwargs
173) -> DataLoader:
174    """Get the MyoFuse dataloader for segmentation of myonuclei.
175
176    Args:
177        path: The folder where the function stores the data.
178        batch_size: The batch size for training.
179        patch_shape: The patch shape to use for training.
180        subset: The cell type. Use 'human' for human primary myotubes or 'mouse' for mouse C2C12 myotubes.
181            The function uses both cell types by default.
182        label_choice: The label type. Use 'semantic' for the class of every nucleus,
183            or 'instances' for the masks of the nuclei.
184        download: Whether to download the data if it is not present.
185        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
186
187    Returns:
188        The DataLoader.
189    """
190    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
191    dataset = get_myofuse_dataset(path, patch_shape, subset, label_choice, download, **ds_kwargs)
192    return torch_em.get_data_loader(dataset, batch_size=batch_size, **loader_kwargs)