torch_em.data.datasets.light_microscopy.plantseg

  1"""This dataset contains confocal and lightsheet microscopy images of plant cells
  2with annotations for cell and nucleus segmentation.
  3
  4The dataset part of the publication https://doi.org/10.7554/eLife.57613.
  5Please cite it if you use this dataset in your research.
  6"""
  7
  8import os
  9from glob import glob
 10from tqdm import tqdm
 11from typing import List, Optional, Tuple, Union
 12
 13from torch.utils.data import Dataset, DataLoader
 14
 15import torch_em
 16
 17from .. import util
 18
 19
 20URLS = {
 21    "root": {
 22        "train": "https://files.de-1.osf.io/v1/resources/9x3g2/providers/osfstorage/?zip=",
 23        "val": "https://files.de-1.osf.io/v1/resources/vs6gb/providers/osfstorage/?zip=",
 24        "test": "https://files.de-1.osf.io/v1/resources/tn4xj/providers/osfstorage/?zip=",
 25    },
 26    "nuclei": {
 27        "train": "https://files.de-1.osf.io/v1/resources/thxzn/providers/osfstorage/?zip=",
 28    },
 29    "ovules": {
 30        "train": "https://files.de-1.osf.io/v1/resources/x9yns/providers/osfstorage/?zip=",
 31        "val": "https://files.de-1.osf.io/v1/resources/xp5uf/providers/osfstorage/?zip=",
 32        "test": "https://files.de-1.osf.io/v1/resources/8jz7e/providers/osfstorage/?zip=",
 33    }
 34}
 35
 36# FIXME somehow the checksums are not reliably, this is a bit weird.
 37CHECKSUMS = {
 38    "root": {
 39        "train": None, "val": None, "test": None
 40        # "train": "f72e9525ff716ef14b70ab1318efd4bf303bbf9e0772bf2981a2db6e22a75794",
 41        # "val": "987280d9a56828c840e508422786431dcc3603e0ba4814aa06e7bf4424efcd9e",
 42        # "test": "ad71b8b9d20effba85fb5e1b42594ae35939d1a0cf905f3403789fc9e6afbc58",
 43    },
 44    "nuclei": {
 45        "train": None
 46        # "train": "9d19ddb61373e2a97effb6cf8bd8baae5f8a50f87024273070903ea8b1160396",
 47    },
 48    "ovules": {
 49        "train": None, "val": None, "test": None
 50        # "train": "70379673f1ab1866df6eb09d5ce11db7d3166d6d15b53a9c8b47376f04bae413",
 51        # "val": "872f516cb76879c30782d9a76d52df95236770a866f75365902c60c37b14fa36",
 52        # "test": "a7272f6ad1d765af6d121e20f436ac4f3609f1a90b1cb2346aa938d8c52800b9",
 53    }
 54}
 55
 56CROPPING_VOLUMES = {
 57    # root (train)
 58    "Movie2_T00006_crop_gt.h5": slice(4, None),
 59    "Movie2_T00008_crop_gt.h5": slice(None, -18),
 60    "Movie2_T00010_crop_gt.h5": slice(None, -32),
 61    "Movie2_T00012_crop_gt.h5": slice(None, -39),
 62    "Movie2_T00014_crop_gt.h5": slice(None, -40),
 63    "Movie2_T00016_crop_gt.h5": slice(None, -42),
 64    # root (test)
 65    "Movie2_T00020_crop_gt.h5": slice(None, -50),
 66    # ovules (train)
 67    "N_487_ds2x.h5": slice(17, None),
 68    "N_535_ds2x.h5": slice(None, -1),
 69    "N_534_ds2x.h5": slice(None, -1),
 70    "N_451_ds2x.h5": slice(None, -1),
 71    "N_425_ds2x.h5": slice(None, -1),
 72    # ovules (val)
 73    "N_420_ds2x.h5": slice(None, -1),
 74}
 75
 76# These root volumes are byte-identical copies of test volumes and are excluded from the train split.
 77DUPLICATE_ROOT_TRAIN_VOLUMES = ("Movie1_t00045_crop_gt.h5", "Movie2_T00010_crop_gt.h5")
 78
 79# In these root volumes the region outside the root carries an instance id instead of the background id 1.
 80ROOT_BACKGROUND_IDS = {
 81    "Movie2_T00000_crop_gt.h5": 286, "Movie3_T00004_crop_gt.h5": 1004, "Movie2_T00020_crop_gt.h5": 411,
 82}
 83
 84# The resolution previous used for the resizing
 85# I have removed this feature since it was not reliable,
 86# but leaving this here for reference
 87# (also implementing resizing would be a good idea,
 88#  but more general and not for each dataset individually)
 89# NATIVE_RESOLUTION = (0.235, 0.075, 0.075)
 90
 91
 92def _fix_inconsistent_volumes(data_path, name, split):
 93    import h5py
 94
 95    file_paths = glob(os.path.join(data_path, "*.h5"))
 96    if name not in ["root", "ovules"] and split not in ["train", "val"]:
 97        return
 98
 99    for vol_path in tqdm(file_paths, desc="Fixing inconsistencies in volumes"):
100        fname = os.path.basename(vol_path)
101        if fname in ROOT_BACKGROUND_IDS:
102            with h5py.File(vol_path, "r+") as f:
103                labels = f["label"][:]
104                labels[labels == ROOT_BACKGROUND_IDS[fname]] = 1
105                f["label"][...] = labels
106
107        if fname not in CROPPING_VOLUMES:
108            continue
109
110        with h5py.File(vol_path, "r+") as f:
111            crop_slices = CROPPING_VOLUMES[fname]
112            for key in ("raw", "label", "label_with_ignore"):
113                if key not in f:
114                    continue
115                ds = f[key]
116                cropped = ds[:][crop_slices]
117                ds.resize(cropped.shape)
118                ds[...] = cropped
119
120
121def get_plantseg_data(path: Union[os.PathLike, str], name: str, split: str, download: bool = False) -> str:
122    """Download the PlantSeg training data.
123
124    Args:
125        path: Filepath to a folder where the downloaded data will be saved.
126        name: The name of the data to load. Either 'root', 'nuclei' or 'ovules'.
127        split: The split to download. Either 'train', 'val' or 'test'.
128        download: Whether to download the data if it is not present.
129
130    Returns:
131        The filepath to the training data.
132    """
133    url = URLS[name][split]
134    checksum = CHECKSUMS[name][split]
135    os.makedirs(path, exist_ok=True)
136    out_path = os.path.join(path, f"{name}_{split}")
137    if os.path.exists(out_path):
138        return out_path
139    tmp_path = os.path.join(path, f"{name}_{split}.zip")
140    util.download_source(tmp_path, url, download, checksum)
141    util.unzip(tmp_path, out_path, remove=True)
142    _fix_inconsistent_volumes(out_path, name, split)
143    return out_path
144
145
146def get_plantseg_paths(
147    path: Union[os.PathLike, str],
148    name: str,
149    split: str,
150    download: bool = False
151) -> List[str]:
152    """Get paths to the PlantSeg data.
153
154    Args:
155        path: Filepath to a folder where the downloaded data will be saved.
156        name: The name of the data to load. Either 'root', 'nuclei' or 'ovules'.
157        split: The split to download. Either 'train', 'val' or 'test'.
158        download: Whether to download the data if it is not present.
159
160    Returns:
161        List of filepaths for the data.
162    """
163    data_path = get_plantseg_data(path, name, split, download)
164    file_paths = sorted(glob(os.path.join(data_path, "*.h5")))
165    if name == "root" and split == "train":
166        file_paths = [p for p in file_paths if os.path.basename(p) not in DUPLICATE_ROOT_TRAIN_VOLUMES]
167    return file_paths
168
169
170def get_plantseg_dataset(
171    path: Union[os.PathLike, str],
172    name: str,
173    split: str,
174    patch_shape: Tuple[int, int, int],
175    download: bool = False,
176    offsets: Optional[List[List[int]]] = None,
177    boundaries: bool = False,
178    binary: bool = False,
179    with_ignore: bool = False,
180    **kwargs,
181) -> Dataset:
182    """Get the PlantSeg dataset for segmenting nuclei or cells.
183
184    Args:
185        path: Filepath to a folder where the downloaded data will be saved.
186        name: The name of the data to load. Either 'root', 'nuclei' or 'ovules'.
187        split: The split to download. Either 'train', 'val' or 'test'.
188        patch_shape: The patch shape to use for training.
189        download: Whether to download the data if it is not present.
190        offsets: Offset values for affinity computation used as target.
191        boundaries: Whether to compute boundaries as the target.
192        binary: Whether to use a binary segmentation target.
193        with_ignore: Whether to load the ovules labels with the unannotated regions marked as -1.
194        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
195
196    Returns:
197       The segmentation dataset.
198    """
199    assert len(patch_shape) == 3
200    if with_ignore and name != "ovules":
201        raise ValueError("Labels with ignore regions are only available for 'ovules'.")
202
203    file_paths = get_plantseg_paths(path, name, split, download)
204
205    kwargs, _ = util.add_instance_label_transform(
206        kwargs, add_binary_target=binary, binary=binary, boundaries=boundaries,
207        offsets=offsets, binary_is_exclusive=False
208    )
209
210    return torch_em.default_segmentation_dataset(
211        raw_paths=file_paths,
212        raw_key="raw",
213        label_paths=file_paths,
214        label_key="label_with_ignore" if with_ignore else "label",
215        patch_shape=patch_shape,
216        **kwargs
217    )
218
219
220def get_plantseg_loader(
221    path: Union[os.PathLike, str],
222    name: str,
223    split: str,
224    patch_shape: Tuple[int, int, int],
225    batch_size: int,
226    download: bool = False,
227    offsets: Optional[List[List[int]]] = None,
228    boundaries: bool = False,
229    binary: bool = False,
230    with_ignore: bool = False,
231    **kwargs,
232) -> DataLoader:
233    """Get the PlantSeg dataloader for segmenting nuclei or cells.
234
235    Args:
236        path: Filepath to a folder where the downloaded data will be saved.
237        name: The name of the data to load. Either 'root', 'nuclei' or 'ovules'.
238        split: The split to download. Either 'train', 'val' or 'test'.
239        patch_shape: The patch shape to use for training.
240        batch_size: The batch size for training.
241        download: Whether to download the data if it is not present.
242        offsets: Offset values for affinity computation used as target.
243        boundaries: Whether to compute boundaries as the target.
244        binary: Whether to use a binary segmentation target.
245        with_ignore: Whether to load the ovules labels with the unannotated regions marked as -1.
246        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
247
248    Returns:
249       The DataLoader.
250    """
251    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
252    dataset = get_plantseg_dataset(
253        path, name, split, patch_shape, download=download, offsets=offsets,
254        boundaries=boundaries, binary=binary, with_ignore=with_ignore, **ds_kwargs
255    )
256    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)