torch_em.data.datasets.light_microscopy.ctc

  1"""The Cell Tracking Challenge contains annotated data for cell segmentation and tracking.
  2We provide the 2d and 3d datasets with segmentation annotations. The 2d datasets are listed in `CTC_2D_DATASETS`,
  3the 3d datasets in `CTC_3D_DATASETS`. See https://celltrackingchallenge.net/2d-datasets/ and
  4https://celltrackingchallenge.net/3d-datasets/ for details on the individual datasets.
  5
  6The segmentation annotations are sparse: only some time points are annotated, and for some of the 3d datasets only
  7individual slices of a time point are annotated. Time points with a fully annotated volume are loaded as 3d data.
  8If a 3d dataset only has slice-wise annotations, the annotated slices are extracted from the raw volumes and loaded
  9as 2d data instead. In addition to the manually curated gold truth annotations ('GT'), the challenge provides
 10silver truth annotations ('ST') that are computationally derived and cover all time points of the 2d and 3d datasets.
 11
 12If you use this data in your research please cite https://doi.org/10.1038/nmeth.4473.
 13"""
 14
 15import os
 16import re
 17from glob import glob
 18from shutil import copyfile
 19from typing import List, Literal, Optional, Tuple, Union
 20
 21import imageio.v3 as imageio
 22from torch.utils.data import Dataset, DataLoader
 23
 24import torch_em
 25
 26from .. import util
 27
 28
 29CTC_2D_DATASETS = [
 30    "BF-C2DL-HSC", "BF-C2DL-MuSC", "DIC-C2DH-HeLa", "Fluo-C2DL-Huh7", "Fluo-C2DL-MSC",
 31    "Fluo-N2DH-GOWT1", "Fluo-N2DH-SIM+", "Fluo-N2DL-HeLa", "PhC-C2DH-U373", "PhC-C2DL-PSC",
 32]
 33"""The names of the 2d datasets."""
 34
 35CTC_3D_DATASETS = [
 36    "Fluo-C3DH-A549", "Fluo-C3DH-A549-SIM", "Fluo-C3DH-H157", "Fluo-C3DL-MDA231", "Fluo-N3DH-CE",
 37    "Fluo-N3DH-CHO", "Fluo-N3DH-SIM+", "Fluo-N3DL-DRO", "Fluo-N3DL-TRIC", "Fluo-N3DL-TRIF",
 38]
 39"""The names of the 3d datasets."""
 40
 41CTC_CHECKSUMS = {
 42    "train": {
 43        "BF-C2DL-HSC": "0aa68ec37a9b06e72a5dfa07d809f56e1775157fb674bb75ff904936149657b1",
 44        "BF-C2DL-MuSC": "ca72b59042809120578a198ba236e5ed3504dd6a122ef969428b7c64f0a5e67d",
 45        "DIC-C2DH-HeLa": "832fed2d05bb7488cf9c51a2994b75f8f3f53b3c3098856211f2d39023c34e1a",
 46        "Fluo-C2DL-Huh7": "1912658c1b3d8b38b314eb658b559e7b39c256917150e9b3dd8bfdc77347617d",
 47        "Fluo-C2DL-MSC": "a083521f0cb673ae02d4957c5e6580c2e021943ef88101f6a2f61b944d671af2",
 48        "Fluo-N2DH-GOWT1": "1a7bd9a7d1d10c4122c7782427b437246fb69cc3322a975485c04e206f64fc2c",
 49        "Fluo-N2DH-SIM+": "3e809148c87ace80c72f563b56c35e0d9448dcdeb461a09c83f61e93f5e40ec8",
 50        "Fluo-N2DL-HeLa": "35dd99d58e071aba0b03880128d920bd1c063783cc280f9531fbdc5be614c82e",
 51        "PhC-C2DH-U373": "b18185c18fce54e8eeb93e4bbb9b201d757add9409bbf2283b8114185a11bc9e",
 52        "PhC-C2DL-PSC": "9d54bb8febc8798934a21bf92e05d92f5e8557c87e28834b2832591cdda78422",
 53        "Fluo-C3DH-A549": "46be7a5403f98070218414e5b71a302f29f697ffcae68f45ecaea5d737076026",
 54        "Fluo-C3DH-A549-SIM": "321bd505854a4e3c9b8255cb903ecc48f43f14abd444448ef6945fc1ec9fa4cf",
 55        "Fluo-C3DH-H157": "9540984397dfb7129b5bff8177b7da11c497708421e6de2713c186b5b15a3c9b",
 56        "Fluo-C3DL-MDA231": "b1044eeaac644f1abfbf91d9b6c97bebfb02f4dee2b2a4aed60d0039a375aa84",
 57        "Fluo-N3DH-CE": "eb3d37cacb3b51d3a427a0c43deaaf1d60cf2eba009f12569e720c2c2fd1e04c",
 58        "Fluo-N3DH-CHO": "48d7e32b6408dddd04f1b6e4153e91f19181b405f4c75610629a25e05d40fe77",
 59        "Fluo-N3DH-SIM+": "f874a297a97ba2b144f2cbe9f06f66f57996251e89b17f7d5ef96a5d99d0fd11",
 60        "Fluo-N3DL-DRO": "57c3039b746116fa068e97f25bb14e8deaf3bd11528d61709c5ae31bc7fc7f11",
 61        "Fluo-N3DL-TRIC": "225f0ed3ba706d9c28101b01bd22ad96974ce2d2334d5c2bf7cddbf2236ec6bb",
 62        # The Fluo-N3DL-TRIF training data is 320 GB large, its checksum has not been computed yet.
 63    },
 64    "test": {
 65        "BF-C2DL-HSC": "fd1c05ec625fd0526c8369d1139babe137e885457eee98c10d957da578d0d5bc",
 66        "BF-C2DL-MuSC": "c5cae259e6090e82a2596967fb54c8a768717c1772398f8546ad1c8df0820450",
 67        "DIC-C2DH-HeLa": "5e5d5f2aa90aef99d750cf03f5c12d799d50b892f98c86950e07a2c5955ac01f",
 68        "Fluo-C2DL-Huh7": "cc7359f8fb6b0c43995365e83ce0116d32f477ac644b2ca02b98bc253e2bcbbe",
 69        "Fluo-C2DL-MSC": "c90b13e603dde52f17801d4f0cadde04ed7f21cc05296b1f0957d92dbfc8ffa6",
 70        "Fluo-N2DH-GOWT1": "c6893ec2d63459de49d4dc21009b04275573403c62cc02e6ee8d0cb1a5068add",
 71        "Fluo-N2DH-SIM+": "c4f257add739b284d02176057814de345dee2ac1a7438e360ccd2df73618db68",
 72        "Fluo-N2DL-HeLa": "45cf3daf05e8495aa2ce0febacca4cf0928fab808c0b14ed2eb7289a819e6bb8",
 73        "PhC-C2DH-U373": "7aa3162e4363a416b259149adc13c9b09cb8aecfe8165eb1428dd534b66bec8a",
 74        "PhC-C2DL-PSC": "8c98ac6203e7490157ceb6aa1131d60a3863001b61fb75e784bc49d47ee264d5",
 75        # The checksums for the test data of the 3d datasets have not been computed yet.
 76    }
 77}
 78"""The checksums of the zip archives for the train and test split of the datasets."""
 79
 80
 81def _get_ctc_url_and_checksum(dataset_name, split):
 82    if split == "train":
 83        _link_to_split = "training-datasets"
 84    else:
 85        _link_to_split = "test-datasets"
 86
 87    url = f"http://data.celltrackingchallenge.net/{_link_to_split}/{dataset_name}.zip"
 88    # The checksum verification is skipped for the archives whose checksum has not been computed yet.
 89    checksum = CTC_CHECKSUMS[split].get(dataset_name)
 90    return url, checksum
 91
 92
 93def get_ctc_segmentation_data(
 94    path: Union[os.PathLike, str], dataset_name: str, split: str, download: bool = False,
 95) -> str:
 96    """Download training data from the Cell Tracking Challenge.
 97
 98    Args:
 99        path: Filepath to a folder where the downloaded data will be saved.
100        dataset_name: Name of the dataset to be downloaded. The available datasets are listed in
101            `CTC_2D_DATASETS` and `CTC_3D_DATASETS`.
102        split: The split to download. Either 'train' or 'test'.
103        download: Whether to download the data if it is not present.
104
105    Returns:
106        The filepath to the training data.
107    """
108    dataset_names = CTC_2D_DATASETS + CTC_3D_DATASETS
109    if dataset_name not in dataset_names:
110        raise ValueError(f"Invalid dataset: {dataset_name}, choose one of {dataset_names}.")
111
112    data_path = os.path.join(path, split, dataset_name)
113
114    if os.path.exists(data_path):
115        return data_path
116
117    os.makedirs(data_path)
118    url, checksum = _get_ctc_url_and_checksum(dataset_name, split)
119    zip_path = os.path.join(path, f"{dataset_name}.zip")
120    util.download_source(zip_path, url, download, checksum=checksum)
121    util.unzip(zip_path, os.path.join(path, split), remove=True)
122
123    return data_path
124
125
126def _parse_label_name(fname):
127    # Annotations of a full time point are named 'man_seg<T>.tif',
128    # annotations of a single slice of a time point are named 'man_seg_<T>_<Z>.tif'.
129    match = re.fullmatch(r"man_seg(\d+)\.tif", fname)
130    if match is not None:
131        return match.group(1), None
132    match = re.fullmatch(r"man_seg_(\d+)_(\d+)\.tif", fname)
133    if match is not None:
134        return match.group(1), int(match.group(2))
135    raise ValueError(f"Unexpected name for a segmentation annotation: {fname}")
136
137
138def _require_gt_image(image_folder, label_path, label_image_folder):
139    fname = os.path.basename(label_path)
140    image_label_path = os.path.join(label_image_folder, fname)
141    if os.path.exists(image_label_path):
142        return image_label_path
143
144    time_point, slice_id = _parse_label_name(fname)
145    image_path = os.path.join(image_folder, f"t{time_point}.tif")
146    assert os.path.exists(image_path), image_path
147
148    if slice_id is None:
149        # Copy over the image corresponding to the fully labeled time point.
150        copyfile(image_path, image_label_path)
151    else:
152        # Extract the labeled slice from the image volume.
153        image = imageio.imread(image_path)
154        imageio.imwrite(image_label_path, image[slice_id])
155
156    return image_label_path
157
158
159def _require_gt_images(data_path, vol_ids, annotation_type="GT", return_files=False):
160    image_paths, label_paths = [], []
161
162    if isinstance(vol_ids, str):
163        vol_ids = [vol_ids]
164
165    # Check whether any of the time points is fully annotated. If so, the fully annotated time points are used.
166    # Otherwise the slice-wise annotations are used.
167    all_label_paths = {
168        vol_id: sorted(glob(os.path.join(data_path, f"{vol_id}_{annotation_type}", "SEG", "*.tif")))
169        for vol_id in vol_ids
170    }
171    assert any(len(paths) > 0 for paths in all_label_paths.values()), f"No annotations found in {data_path}."
172    use_slices = all(
173        _parse_label_name(os.path.basename(p))[1] is not None for paths in all_label_paths.values() for p in paths
174    )
175
176    for vol_id in vol_ids:
177        image_folder = os.path.join(data_path, vol_id)
178        assert os.path.exists(image_folder), f"Cannot find volume id, {vol_id} in {data_path}."
179
180        label_folder = os.path.join(data_path, f"{vol_id}_{annotation_type}", "SEG")
181
182        # Copy over the images corresponding to the labeled frames.
183        label_image_folder = os.path.join(data_path, f"{vol_id}_{annotation_type}", "IM")
184        os.makedirs(label_image_folder, exist_ok=True)
185
186        this_label_paths = [
187            p for p in all_label_paths[vol_id] if (_parse_label_name(os.path.basename(p))[1] is not None) == use_slices
188        ]
189        this_image_paths = [_require_gt_image(image_folder, p, label_image_folder) for p in this_label_paths]
190
191        if return_files:
192            image_paths.extend(this_image_paths)
193            label_paths.extend(this_label_paths)
194        else:
195            image_paths.append(label_image_folder)
196            label_paths.append(label_folder)
197
198    return image_paths, label_paths
199
200
201def get_ctc_segmentation_paths(
202    path: Union[os.PathLike, str],
203    dataset_name: str,
204    split: str = "train",
205    vol_id: Optional[int] = None,
206    download: bool = False,
207    annotation_type: Literal["GT", "ST"] = "GT",
208) -> Tuple[List[str], List[str]]:
209    """Get paths to the Cell Tracking Challenge data.
210
211    For the 2d datasets this returns the folders with the images of the annotated time points and the folders
212    with the corresponding annotations. For the 3d datasets this returns the filepaths of the annotated volumes
213    (or the annotated slices for datasets that only have slice-wise annotations) and of the corresponding annotations.
214
215    Args:
216        path: Filepath to a folder where the downloaded data will be saved.
217        dataset_name: Name of the dataset to be downloaded. The available datasets are listed in
218            `CTC_2D_DATASETS` and `CTC_3D_DATASETS`.
219        split: The split to download. Currently only supports 'train'.
220        vol_id: The train id to load.
221        download: Whether to download the data if it is not present.
222        annotation_type: The type of annotations to load. Either 'GT' for the manually curated gold truth
223            annotations or 'ST' for the computationally derived silver truth annotations.
224
225    Returns:
226        Filepaths to the image data.
227        Filepaths to the label data.
228    """
229    data_path = get_ctc_segmentation_data(path, dataset_name, split, download)
230
231    if vol_id is None:
232        vol_ids = glob(os.path.join(data_path, "*_GT"))
233        vol_ids = [os.path.basename(vol_id) for vol_id in vol_ids]
234        vol_ids = sorted(vol_id[:-len("_GT")] for vol_id in vol_ids)
235    else:
236        vol_ids = vol_id
237
238    assert annotation_type in ("GT", "ST"), f"Invalid annotation type: {annotation_type}, choose 'GT' or 'ST'."
239    return_files = dataset_name in CTC_3D_DATASETS
240    image_path, label_path = _require_gt_images(data_path, vol_ids, annotation_type, return_files=return_files)
241    return image_path, label_path
242
243
244def get_ctc_segmentation_dataset(
245    path: Union[os.PathLike, str],
246    dataset_name: str,
247    patch_shape: Tuple[int, ...],
248    split: str = "train",
249    vol_id: Optional[int] = None,
250    download: bool = False,
251    annotation_type: Literal["GT", "ST"] = "GT",
252    **kwargs,
253) -> Dataset:
254    """Get the CTC dataset for cell segmentation.
255
256    Args:
257        path: Filepath to a folder where the downloaded data will be saved.
258        dataset_name: Name of the dataset to be downloaded. The available datasets are listed in
259            `CTC_2D_DATASETS` and `CTC_3D_DATASETS`.
260        patch_shape: The patch shape to use for training. Use a 2d patch shape (with a leading 1) for the 2d datasets
261            and a 3d patch shape for the 3d datasets (a 2d patch shape for the 3d datasets with slice-wise annotations).
262        split: The split to download. Currently only supports 'train'.
263        vol_id: The train id to load.
264        download: Whether to download the data if it is not present.
265        annotation_type: The type of annotations to load. Either 'GT' for the manually curated gold truth
266            annotations or 'ST' for the computationally derived silver truth annotations.
267        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
268
269    Returns:
270       The segmentation dataset.
271    """
272    assert split in ["train"]
273
274    image_path, label_path = get_ctc_segmentation_paths(path, dataset_name, split, vol_id, download, annotation_type)
275
276    if dataset_name in CTC_3D_DATASETS:
277        # The data is loaded as 3d if it has fully annotated time points and as 2d if it only has annotated slices.
278        ndim = 2 if _parse_label_name(os.path.basename(label_path[0]))[1] is not None else 3
279        kwargs = util.update_kwargs(kwargs, "ndim", ndim)
280        raw_key, label_key = None, None
281    else:
282        kwargs = util.update_kwargs(kwargs, "ndim", 2)
283        raw_key, label_key = "*.tif", "*.tif"
284
285    return torch_em.default_segmentation_dataset(
286        raw_paths=image_path,
287        raw_key=raw_key,
288        label_paths=label_path,
289        label_key=label_key,
290        patch_shape=patch_shape,
291        is_seg_dataset=True,
292        **kwargs
293    )
294
295
296def get_ctc_segmentation_loader(
297    path: Union[os.PathLike, str],
298    dataset_name: str,
299    patch_shape: Tuple[int, ...],
300    batch_size: int,
301    split: str = "train",
302    vol_id: Optional[int] = None,
303    download: bool = False,
304    annotation_type: Literal["GT", "ST"] = "GT",
305    **kwargs,
306) -> DataLoader:
307    """Get the CTC dataloader for cell segmentation.
308
309    Args:
310        path: Filepath to a folder where the downloaded data will be saved.
311        dataset_name: Name of the dataset to be downloaded. The available datasets are listed in
312            `CTC_2D_DATASETS` and `CTC_3D_DATASETS`.
313        patch_shape: The patch shape to use for training. Use a 2d patch shape (with a leading 1) for the 2d datasets
314            and a 3d patch shape for the 3d datasets (a 2d patch shape for the 3d datasets with slice-wise annotations).
315        batch_size: The batch size for training.
316        split: The split to download. Currently only supports 'train'.
317        vol_id: The train id to load.
318        download: Whether to download the data if it is not present.
319        annotation_type: The type of annotations to load. Either 'GT' for the manually curated gold truth
320            annotations or 'ST' for the computationally derived silver truth annotations.
321        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
322
323    Returns:
324       The DataLoader.
325    """
326    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
327    dataset = get_ctc_segmentation_dataset(
328        path, dataset_name, patch_shape, split, vol_id, download, annotation_type, **ds_kwargs
329    )
330    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)