torch_em.data.datasets.light_microscopy.sperm_scd

  1"""The SCD Sperm dataset contains annotations for sperm instance segmentation in bright-field microscopy
  2images of the sperm chromatin dispersion (SCD) assay, which is used to assess sperm DNA fragmentation.
  3
  4The images (1920x1080, RGB) were acquired from 40 coded slides (L01 to L40) and every sperm cell is annotated
  5with a polygon and one of two classes: fragmented (halo-less) or non-fragmented (with halo). The official splits
  6are defined per slide: L01-L32 for 'train', L33-L36 for 'val' and L37-L40 for 'test'. This module rasterizes
  7the COCO polygons into instance labels (one id per sperm, starting from 1) or into class labels
  8(1 = fragmented, 2 = non-fragmented). If polygons overlap, the smaller instance is painted last and wins.
  9
 10NOTE: The release only ships the images of a subset of the slides (about 1,100 of the 2,000 images described in its
 11README), while the annotations are provided for all slides. This module only uses annotated images that are
 12present on disk, which are discovered from the files, and it ignores the images without annotated sperm.
 13
 14The dataset is located at https://doi.org/10.5281/zenodo.21628868 and is released under a CC-BY-4.0 license.
 15Please cite it if you use this dataset for your research.
 16"""
 17
 18import os
 19import json
 20from glob import glob
 21from natsort import natsorted
 22from typing import Union, Tuple, List, Literal
 23
 24import numpy as np
 25from tqdm import tqdm
 26
 27from torch.utils.data import Dataset, DataLoader
 28
 29import torch_em
 30
 31from .. import util
 32
 33
 34URL = "https://zenodo.org/api/records/21628868/files/SCD_Microscopy_YOLOv8seg_Repository.zip/content"
 35CHECKSUM = "6e8d9e03f9037c7316113572497fc5ac61ddc2dc63f504b4af91a01c0ca1afeb"
 36
 37SPLITS = ["train", "val", "test"]
 38LABEL_TYPES = ["instances", "classes"]
 39
 40
 41def _slide_splits(root):
 42    slide_to_split = {}
 43    for split in SPLITS:
 44        for label_path in glob(os.path.join(root, "dataset", "labels_split", split, "*.txt")):
 45            slide_to_split[os.path.basename(label_path)[:3]] = split
 46    return slide_to_split
 47
 48
 49def _find_images(root):
 50    image_paths = {}
 51    for pattern in (("dataset", "images", "*.png"), ("dataset", "images_split", "*", "*.png")):
 52        for image_path in glob(os.path.join(root, *pattern)):
 53            image_paths.setdefault(os.path.basename(image_path), image_path)
 54    return image_paths
 55
 56
 57def _rasterize(annotations, shape):
 58    from skimage.draw import polygon
 59
 60    instances = np.zeros(shape, dtype="uint16")
 61    classes = np.zeros(shape, dtype="uint8")
 62    for instance_id, annotation in enumerate(sorted(annotations, key=lambda a: -a["area"]), start=1):
 63        for part in annotation["segmentation"]:
 64            coords = np.asarray(part, dtype="float64").reshape(-1, 2)
 65            rr, cc = polygon(coords[:, 1], coords[:, 0], shape=shape)
 66            instances[rr, cc] = instance_id
 67            classes[rr, cc] = annotation["category_id"]
 68    return instances, classes
 69
 70
 71def _write_tif(path, array):
 72    import tifffile
 73
 74    tmp_path = f"{path}.{os.getpid()}.incomplete.tif"
 75    tifffile.imwrite(tmp_path, array, compression="zlib")
 76    os.replace(tmp_path, path)
 77
 78
 79def _preprocess_labels(root, label_root):
 80    done_marker = os.path.join(label_root, "done")
 81    if os.path.exists(done_marker):
 82        return
 83
 84    image_paths = _find_images(root)
 85
 86    per_image = {}
 87    for json_path in natsorted(glob(os.path.join(root, "annotations", "coco_json", "*.json"))):
 88        with open(json_path) as f:
 89            coco = json.load(f)
 90        file_names = {image["id"]: image for image in coco["images"]}
 91        for annotation in coco["annotations"]:
 92            image = file_names[annotation["image_id"]]
 93            per_image.setdefault(image["file_name"], (image, []))[1].append(annotation)
 94
 95    for label_type in LABEL_TYPES:
 96        os.makedirs(os.path.join(label_root, label_type), exist_ok=True)
 97
 98    for file_name, (image, annotations) in tqdm(sorted(per_image.items()), desc="Preprocess SCD Sperm"):
 99        if file_name not in image_paths:
100            continue
101
102        stem = os.path.splitext(file_name)[0]
103        out_paths = [os.path.join(label_root, label_type, f"{stem}.tif") for label_type in LABEL_TYPES]
104        if all(os.path.exists(p) for p in out_paths):
105            continue
106
107        instances, classes = _rasterize(annotations, (image["height"], image["width"]))
108        for out_path, array in zip(out_paths, (instances, classes)):
109            _write_tif(out_path, array)
110
111    with open(done_marker, "w"):
112        pass
113
114
115def get_sperm_scd_data(path: Union[os.PathLike, str], download: bool = False) -> str:
116    """Download the SCD Sperm dataset and rasterize its annotations.
117
118    Args:
119        path: Filepath to a folder where the data is downloaded for further processing.
120        download: Whether to download the data if it is not present.
121
122    Returns:
123        Filepath to the folder with the extracted data.
124    """
125    root = os.path.join(path, "SCD_Microscopy_YOLOv8seg_Repository")
126    if not os.path.exists(root):
127        os.makedirs(path, exist_ok=True)
128        zip_path = os.path.join(path, "SCD_Microscopy_YOLOv8seg_Repository.zip")
129        util.download_source(path=zip_path, url=URL, download=download, checksum=CHECKSUM)
130        util.unzip(zip_path=zip_path, dst=path, remove=False)
131
132    _preprocess_labels(root, os.path.join(path, "labels"))
133    return root
134
135
136def get_sperm_scd_paths(
137    path: Union[os.PathLike, str],
138    split: Literal["train", "val", "test"],
139    label_type: Literal["instances", "classes"] = "instances",
140    download: bool = False,
141) -> Tuple[List[str], List[str]]:
142    """Get paths to the SCD Sperm data.
143
144    Args:
145        path: Filepath to a folder where the data is downloaded for further processing.
146        split: The choice of data split. One of 'train', 'val' or 'test'.
147        label_type: The choice of labels. Either 'instances' (one id per sperm) or 'classes'
148            (1 = fragmented, 2 = non-fragmented).
149        download: Whether to download the data if it is not present.
150
151    Returns:
152        List of filepaths for the image data.
153        List of filepaths for the label data.
154    """
155    if split not in SPLITS:
156        raise ValueError(f"'{split}' is not a valid split. Choose one of {SPLITS}.")
157    if label_type not in LABEL_TYPES:
158        raise ValueError(f"'{label_type}' is not a valid label type. Choose one of {LABEL_TYPES}.")
159
160    root = get_sperm_scd_data(path, download)
161    slide_to_split = _slide_splits(root)
162    image_paths = _find_images(root)
163
164    raw_paths, label_paths = [], []
165    for label_path in natsorted(glob(os.path.join(path, "labels", label_type, "*.tif"))):
166        stem = os.path.splitext(os.path.basename(label_path))[0]
167        if slide_to_split.get(stem[:3]) == split and f"{stem}.png" in image_paths:
168            raw_paths.append(image_paths[f"{stem}.png"])
169            label_paths.append(label_path)
170
171    assert len(raw_paths) == len(label_paths) and len(raw_paths) > 0
172    return raw_paths, label_paths
173
174
175def get_sperm_scd_dataset(
176    path: Union[os.PathLike, str],
177    patch_shape: Tuple[int, int],
178    split: Literal["train", "val", "test"],
179    label_type: Literal["instances", "classes"] = "instances",
180    download: bool = False,
181    **kwargs
182) -> Dataset:
183    """Get the SCD Sperm dataset for sperm instance segmentation.
184
185    Args:
186        path: Filepath to a folder where the data is downloaded for further processing.
187        patch_shape: The patch shape to use for training.
188        split: The choice of data split. One of 'train', 'val' or 'test'.
189        label_type: The choice of labels. Either 'instances' (one id per sperm) or 'classes'
190            (1 = fragmented, 2 = non-fragmented).
191        download: Whether to download the data if it is not present.
192        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
193
194    Returns:
195        The segmentation dataset.
196    """
197    raw_paths, label_paths = get_sperm_scd_paths(path, split, label_type, download)
198
199    return torch_em.default_segmentation_dataset(
200        raw_paths=raw_paths,
201        raw_key=None,
202        label_paths=label_paths,
203        label_key=None,
204        is_seg_dataset=False,
205        patch_shape=patch_shape,
206        **kwargs
207    )
208
209
210def get_sperm_scd_loader(
211    path: Union[os.PathLike, str],
212    batch_size: int,
213    patch_shape: Tuple[int, int],
214    split: Literal["train", "val", "test"],
215    label_type: Literal["instances", "classes"] = "instances",
216    download: bool = False,
217    **kwargs
218) -> DataLoader:
219    """Get the SCD Sperm dataloader for sperm instance segmentation.
220
221    Args:
222        path: Filepath to a folder where the data is downloaded for further processing.
223        batch_size: The batch size for training.
224        patch_shape: The patch shape to use for training.
225        split: The choice of data split. One of 'train', 'val' or 'test'.
226        label_type: The choice of labels. Either 'instances' (one id per sperm) or 'classes'
227            (1 = fragmented, 2 = non-fragmented).
228        download: Whether to download the data if it is not present.
229        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
230
231    Returns:
232        The DataLoader.
233    """
234    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
235    dataset = get_sperm_scd_dataset(path, patch_shape, split, label_type, download, **ds_kwargs)
236    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)