torch_em.data.datasets.light_microscopy.pnas_arabidopsis

  1"""The PNAS Arabidopsis dataset contains cell segmentation in confocal microscopy images of
  2arabidopsis plantlets.
  3
  4NOTE: There is tracking information available for this data.
  5
  6This dataset is from the publication https://doi.org/10.1073/pnas.1616768113.
  7Please cite it if you use this dataset for your research.
  8"""
  9
 10import os
 11import shutil
 12from glob import glob
 13from tqdm import tqdm
 14from pathlib import Path
 15from natsort import natsorted
 16from typing import Union, Tuple, List, Optional, Sequence
 17
 18import imageio.v3 as imageio
 19
 20from torch.utils.data import Dataset, DataLoader
 21
 22import torch_em
 23
 24from .. import util
 25
 26
 27URL = "https://www.repository.cam.ac.uk/bitstream/handle/1810/262530/PNAS.zip?sequence=4&isAllowed=y"
 28CHECKSUM = "39341398389baf6d93c3f652b7e2e8aedc5579c29dfaf2b82b41ebfc3caa05c4"
 29
 30
 31def get_pnas_arabidopsis_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 32    """Download the PNAS Arabidopsis dataset.
 33
 34    Args:
 35        path: Filepath to a folder where the data is downloaded for further processing.
 36        download: Whether to download the data if it is not present.
 37
 38    Returns:
 39        Filepath where the data is downloaded and pre-processed.
 40    """
 41    data_dir = os.path.join(path, "data")
 42    if os.path.exists(data_dir):
 43        return data_dir
 44
 45    os.makedirs(data_dir)
 46
 47    zip_path = os.path.join(path, "PNAS.zip")
 48    util.download_source(path=zip_path, url=URL, download=download, checksum=CHECKSUM)
 49    util.unzip(zip_path=zip_path, dst=data_dir)
 50
 51    # Convert the data to h5 (It's hard to keep a track of filenames as they are not completely consistent)
 52    import h5py
 53
 54    raw_paths = natsorted(glob(os.path.join(data_dir, "PNAS", "plant*", "processed_tiffs", "*trim-acylYFP.tif")))
 55    for rpath in tqdm(raw_paths, desc="Preprocessing images"):
 56        # Let's find the label.
 57        label_path = rpath.replace("processed_tiffs", "segmentation_tiffs")
 58        label_path = glob(label_path.replace(".tif", "*.tif"))
 59
 60        if len(label_path) != 1:
 61            print(f"It seems like there are no matching labels for '{os.path.basename(rpath)}'.")
 62            continue
 63
 64        label_path = label_path[0]
 65
 66        raw = imageio.imread(rpath)
 67        labels = imageio.imread(label_path)
 68
 69        # Store both image and corresponding labels in a h5 file.
 70        vol_path = os.path.join(data_dir, Path(os.path.basename(rpath)).with_suffix(".h5"))
 71        with h5py.File(vol_path, "w") as f:
 72            f.create_dataset("raw", data=raw, dtype=raw.dtype, compression="gzip")
 73            f.create_dataset("labels", data=labels, dtype=labels.dtype, compression="gzip")
 74
 75    # Remove old data folder
 76    shutil.rmtree(os.path.join(path, "data", "PNAS"))
 77
 78    return data_dir
 79
 80
 81def get_pnas_arabidopsis_paths(
 82    path: Union[os.PathLike, str], plants: Optional[Sequence[str]] = None, download: bool = False
 83) -> List[str]:
 84    """Get paths to the PNAS Arabidopsis data.
 85
 86    Args:
 87        path: Filepath to a folder where the data is downloaded for further processing.
 88        plants: The plants to restrict to, e.g. ["plant4"]. By default all six plants are used.
 89        download: Whether to download the data if it is not present.
 90
 91    Returns:
 92        List of filepaths for the volumetric data.
 93    """
 94    data_dir = get_pnas_arabidopsis_data(path, download)
 95    volume_paths = sorted(glob(os.path.join(data_dir, "*.h5")))
 96    if plants is not None:
 97        volume_paths = [p for p in volume_paths if os.path.basename(p).split("_")[1] in plants]
 98    return volume_paths
 99
100
101def get_pnas_arabidopsis_dataset(
102    path: Union[os.PathLike, str],
103    patch_shape: Tuple[int, ...],
104    plants: Optional[Sequence[str]] = None,
105    download: bool = False,
106    **kwargs
107) -> Dataset:
108    """Get the PNAS Arabidopsis dataset for cell segmentation.
109
110    Args:
111        path: Filepath to a folder where the data is downloaded for further processing.
112        patch_shape: The patch shape to use for training.
113        plants: The plants to restrict to, e.g. ["plant4"]. By default all six plants are used.
114        download: Whether to download the data if it is not present.
115        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
116
117    Returns:
118        The segmentation dataset.
119    """
120    volume_paths = get_pnas_arabidopsis_paths(path, plants, download)
121
122    return torch_em.default_segmentation_dataset(
123        raw_paths=volume_paths,
124        raw_key="raw",
125        label_paths=volume_paths,
126        label_key="labels",
127        patch_shape=patch_shape,
128        is_seg_dataset=True,
129        **kwargs
130    )
131
132
133def get_pnas_arabidopsis_loader(
134    path: Union[os.PathLike, str],
135    batch_size: int,
136    patch_shape: Tuple[int, ...],
137    plants: Optional[Sequence[str]] = None,
138    download: bool = False,
139    **kwargs
140) -> DataLoader:
141    """Get the PNAS Arabidopsis dataset for cell segmentation.
142
143    Args:
144        path: Filepath to a folder where the data is downloaded for further processing.
145        batch_size: The batch size for training.
146        patch_shape: The patch shape to use for training.
147        plants: The plants to restrict to, e.g. ["plant4"]. By default all six plants are used.
148        download: Whether to download the data if it is not present.
149        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
150
151    Returns:
152        The segmentation dataset.
153    """
154    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
155    dataset = get_pnas_arabidopsis_dataset(path, patch_shape, plants, download, **ds_kwargs)
156    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)