torch_em.data.datasets.light_microscopy.stomataquant

  1"""The StomataQuant dataset contains annotations for stomata, stomatal pore and pavement cell segmentation
  2in bright-field microscopy images of leaf epidermis.
  3
  4Two segmentation tasks are available, selected with the 'task' argument:
  5- 'pores': stomata (class 0) and stomatal pores (class 1), from 500 development images.
  6- 'pavement_cells': stomata (class 0) and pavement cells (class 1), from 613 development images.
  7
  8The development images are split into 'train' and 'val'. The independent test images of the publication
  9(45 per task, from other species and imaging conditions) are available as the 'test' split. The annotations are
 10distributed as YOLO polygons (class id and normalized polygon coordinates), which are rasterized to label images
 11by this module. The class ids are not documented by the authors. The order above was inferred from the data:
 12pores are small polygons lying inside stomata polygons, and pavement cells are the more numerous class.
 13
 14Label images can be one of two types, selected with the 'label_type' argument:
 15- 'instances': every polygon gets its own id. Where polygons overlap, stomata are drawn over pavement cells and
 16  pores are drawn over stomata (see `INSTANCE_DRAW_ORDER`).
 17- 'semantic': the background is 0 and every polygon is labeled with its class id plus one, using the same
 18  drawing order.
 19
 20The images differ in size and some are stored as grayscale, palette or RGBA images. These are converted to RGB.
 21The large detection data (Supplementary Dataset S1-1) contains bounding boxes only and is not used.
 22
 23The data is located at https://doi.org/10.5281/zenodo.18934358 and released under a CC-BY-4.0 license.
 24This dataset is from the publication https://doi.org/10.1093/jpe/rtag063.
 25Please cite it if you use this dataset in your research.
 26"""
 27
 28import os
 29import uuid
 30from glob import glob
 31from tqdm import tqdm
 32from natsort import natsorted
 33from concurrent import futures
 34from typing import Union, Tuple, Optional, Literal, List
 35
 36from torch.utils.data import Dataset, DataLoader
 37
 38import torch_em
 39
 40from .. import util
 41
 42
 43BASE_URL = "https://zenodo.org/api/records/18934358/files"
 44
 45FILES = {
 46    "pores": "Supplementary Dataset S1-2_Stomata_and_pores_segmentation_model.zip",
 47    "pavement_cells": "Supplementary Dataset S1-3_Stomata_and_pavement_cells_segmentation_model.zip",
 48    "test": "Supplementary Dataset S2.zip",
 49}
 50
 51CHECKSUMS = {
 52    "pores": "d6c06377cf21e5f9c42fc2f65df6cd1f5c18b9c36e4fd3edca5eef4fac618eef",
 53    "pavement_cells": "8db4ccb60ba4f8153a7f915b6691fa9841eb70ed9462f2cf5e0387b39d02eb7f",
 54    "test": "4378b7a6ffd125e7ef036b7ca9baa8aa03ccd771c6009407e9450b6eeb0c705c",
 55}
 56
 57TASKS = ("pores", "pavement_cells")
 58SPLITS = ("train", "val", "test")
 59LABEL_TYPES = ("instances", "semantic")
 60
 61TEST_FOLDERS = {
 62    "pores": "Test_stomata_and_pores_segmentation_model",
 63    "pavement_cells": "Test_stomata_and_pavement_cells_segmentation_model",
 64}
 65
 66INSTANCE_DRAW_ORDER = {"pores": (0, 1), "pavement_cells": (1, 0)}
 67"""The order in which the classes are drawn for each task, later classes overwrite earlier ones."""
 68
 69
 70def _read_polygons(txt_path):
 71    import numpy as np
 72
 73    polygons = {}
 74    with open(txt_path) as f:
 75        for line in f:
 76            tokens = line.split()
 77            if len(tokens) < 7 or (len(tokens) - 1) % 2 != 0:
 78                continue
 79            polygons.setdefault(int(tokens[0]), []).append(np.array(tokens[1:], dtype="float64").reshape(-1, 2))
 80    return polygons
 81
 82
 83def _write_atomic(path, array):
 84    import imageio.v3 as imageio
 85
 86    extension = os.path.splitext(path)[1]
 87    tmp_path = f"{os.path.splitext(path)[0]}.{uuid.uuid4().hex}.incomplete{extension}"
 88    imageio.imwrite(tmp_path, array)
 89    os.replace(tmp_path, path)
 90
 91
 92def _process_item(image_path, txt_path, rgb_path, label_path, task, label_type):
 93    import numpy as np
 94    from PIL import Image, ImageDraw
 95
 96    with Image.open(image_path) as image:
 97        width, height = image.size
 98        needs_conversion = image.mode != "RGB"
 99
100    if needs_conversion and not os.path.exists(rgb_path):
101        with Image.open(image_path) as image:
102            _write_atomic(rgb_path, np.asarray(image.convert("RGB")))
103
104    if os.path.exists(label_path):
105        return
106
107    polygons = _read_polygons(txt_path)
108    # skimage.draw.polygon needs about 0.5 s per polygon on these large images, the PIL scanline fill takes ms.
109    canvas = Image.new("I", (width, height), 0)
110    draw = ImageDraw.Draw(canvas)
111    next_id = 1
112    for class_id in INSTANCE_DRAW_ORDER[task]:
113        for points in polygons.get(class_id, []):
114            vertices = [(x * width, y * height) for x, y in points]
115            draw.polygon(vertices, fill=next_id if label_type == "instances" else class_id + 1)
116            next_id += 1
117
118    _write_atomic(label_path, np.asarray(canvas).astype("uint16"))
119
120
121def _list_items(data_dir, task, split):
122    if split == "test":
123        folder = os.path.join(data_dir, "test", TEST_FOLDERS[task])
124        image_paths = natsorted(glob(os.path.join(folder, "imgs", "*")))
125        label_dir = os.path.join(folder, "groundtruth_labels")
126    else:
127        folder = os.path.join(data_dir, task)
128        image_paths = natsorted(glob(os.path.join(folder, "images", split, "*")))
129        label_dir = os.path.join(folder, "labels", split)
130
131    items = []
132    for image_path in image_paths:
133        stem = os.path.splitext(os.path.basename(image_path))[0]
134        txt_path = os.path.join(label_dir, f"{stem}.txt")
135        if os.path.exists(txt_path):
136            items.append((image_path, txt_path, stem))
137
138    return items
139
140
141def get_stomataquant_data(
142    path: Union[os.PathLike, str], task: Literal["pores", "pavement_cells"], split: Literal["train", "val", "test"],
143    download: bool = False,
144) -> str:
145    """Download the StomataQuant dataset.
146
147    Args:
148        path: Filepath to a folder where the downloaded data will be saved.
149        task: The segmentation task. Either 'pores' or 'pavement_cells'.
150        split: The data split. One of 'train', 'val' or 'test'.
151        download: Whether to download the data if it is not present.
152
153    Returns:
154        Filepath where the data is downloaded.
155    """
156    if task not in TASKS:
157        raise ValueError(f"'{task}' is not a valid task. Choose one of {TASKS}.")
158    if split not in SPLITS:
159        raise ValueError(f"'{split}' is not a valid split. Choose one of {SPLITS}.")
160
161    archive = "test" if split == "test" else task
162    dst = os.path.join(path, archive)
163    if os.path.exists(dst) and os.listdir(dst):
164        return path
165
166    os.makedirs(path, exist_ok=True)
167    zip_path = os.path.join(path, f"{archive}.zip")
168    url = f"{BASE_URL}/{FILES[archive].replace(' ', '%20')}/content"
169    util.download_source(path=zip_path, url=url, download=download, checksum=CHECKSUMS[archive])
170    util.unzip(zip_path=zip_path, dst=dst, remove=False)
171
172    return path
173
174
175def get_stomataquant_paths(
176    path: Union[os.PathLike, str],
177    task: Literal["pores", "pavement_cells"],
178    split: Literal["train", "val", "test"],
179    label_type: Literal["instances", "semantic"] = "instances",
180    download: bool = False,
181) -> Tuple[List[str], List[str]]:
182    """Get paths to the StomataQuant data.
183
184    Args:
185        path: Filepath to a folder where the downloaded data will be saved.
186        task: The segmentation task. Either 'pores' or 'pavement_cells'.
187        split: The data split. One of 'train', 'val' or 'test'.
188        label_type: The type of label image. Either 'instances' or 'semantic'.
189        download: Whether to download the data if it is not present.
190
191    Returns:
192        List of filepaths for the image data.
193        List of filepaths for the label data.
194    """
195    if label_type not in LABEL_TYPES:
196        raise ValueError(f"'{label_type}' is not a valid label type. Choose one of {LABEL_TYPES}.")
197
198    data_dir = get_stomataquant_data(path, task, split, download)
199    items = _list_items(data_dir, task, split)
200    assert len(items) > 0, f"No images with annotations were found for task '{task}' and split '{split}'."
201
202    rgb_dir = os.path.join(path, "images_rgb", task, split)
203    label_dir = os.path.join(path, "labels", task, label_type, split)
204    os.makedirs(rgb_dir, exist_ok=True)
205    os.makedirs(label_dir, exist_ok=True)
206
207    jobs = [
208        (image_path, txt_path, os.path.join(rgb_dir, f"{stem}.png"), os.path.join(label_dir, f"{stem}.tif"))
209        for image_path, txt_path, stem in items
210    ]
211    todo = [job for job in jobs if not os.path.exists(job[3])]
212    if todo:
213        with futures.ThreadPoolExecutor(min(8, os.cpu_count() or 1)) as pool:
214            tasks = [pool.submit(_process_item, *job, task, label_type) for job in todo]
215            for job in tqdm(futures.as_completed(tasks), total=len(tasks), desc="Preprocess StomataQuant"):
216                job.result()
217
218    image_paths = [job[2] if os.path.exists(job[2]) else job[0] for job in jobs]
219    label_paths = [job[3] for job in jobs]
220    return image_paths, label_paths
221
222
223def get_stomataquant_dataset(
224    path: Union[os.PathLike, str],
225    patch_shape: Tuple[int, int],
226    task: Literal["pores", "pavement_cells"],
227    split: Literal["train", "val", "test"],
228    label_type: Literal["instances", "semantic"] = "instances",
229    offsets: Optional[List[List[int]]] = None,
230    boundaries: bool = False,
231    binary: bool = False,
232    resize_inputs: bool = False,
233    download: bool = False,
234    **kwargs
235) -> Dataset:
236    """Get the StomataQuant dataset for stomata, pore and pavement cell segmentation.
237
238    Args:
239        path: Filepath to a folder where the downloaded data will be saved.
240        patch_shape: The patch shape to use for training.
241        task: The segmentation task. Either 'pores' or 'pavement_cells'.
242        split: The data split. One of 'train', 'val' or 'test'.
243        label_type: The type of label image. Either 'instances' or 'semantic'.
244        offsets: Offset values for affinity computation used as target.
245        boundaries: Whether to compute boundaries as the target.
246        binary: Whether to use a binary segmentation target.
247        resize_inputs: Whether to resize the inputs to the patch shape. The images differ in size.
248        download: Whether to download the data if it is not present.
249        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
250
251    Returns:
252        The segmentation dataset.
253    """
254    image_paths, label_paths = get_stomataquant_paths(path, task, split, label_type, download)
255
256    if resize_inputs:
257        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": True}
258        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
259            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
260        )
261
262    if label_type == "instances":
263        kwargs, _ = util.add_instance_label_transform(
264            kwargs, add_binary_target=True, offsets=offsets, boundaries=boundaries, binary=binary
265        )
266
267    return torch_em.default_segmentation_dataset(
268        raw_paths=image_paths,
269        raw_key=None,
270        label_paths=label_paths,
271        label_key=None,
272        patch_shape=patch_shape,
273        is_seg_dataset=False,
274        **kwargs
275    )
276
277
278def get_stomataquant_loader(
279    path: Union[os.PathLike, str],
280    batch_size: int,
281    patch_shape: Tuple[int, int],
282    task: Literal["pores", "pavement_cells"],
283    split: Literal["train", "val", "test"],
284    label_type: Literal["instances", "semantic"] = "instances",
285    offsets: Optional[List[List[int]]] = None,
286    boundaries: bool = False,
287    binary: bool = False,
288    resize_inputs: bool = False,
289    download: bool = False,
290    **kwargs
291) -> DataLoader:
292    """Get the StomataQuant dataloader for stomata, pore and pavement cell segmentation.
293
294    Args:
295        path: Filepath to a folder where the downloaded data will be saved.
296        batch_size: The batch size for training.
297        patch_shape: The patch shape to use for training.
298        task: The segmentation task. Either 'pores' or 'pavement_cells'.
299        split: The data split. One of 'train', 'val' or 'test'.
300        label_type: The type of label image. Either 'instances' or 'semantic'.
301        offsets: Offset values for affinity computation used as target.
302        boundaries: Whether to compute boundaries as the target.
303        binary: Whether to use a binary segmentation target.
304        resize_inputs: Whether to resize the inputs to the patch shape. The images differ in size.
305        download: Whether to download the data if it is not present.
306        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
307
308    Returns:
309        The DataLoader.
310    """
311    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
312    dataset = get_stomataquant_dataset(
313        path, patch_shape, task, split, label_type, offsets, boundaries, binary, resize_inputs, download, **ds_kwargs
314    )
315    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)