torch_em.data.datasets.light_microscopy.svia

  1"""The SVIA dataset contains annotations for sperm segmentation in
  2bright-field microscopy videos of human semen.
  3
  4SVIA stands for Sperm Videos and Images Analysis. The deposit calls the dataset MIaMIA-SVDS. This
  5loader uses Subset-B, which holds a mask per sperm for 451 frames of ten videos, so 25966 objects.
  6
  7NOTE: The dataset holds two more label types, which this loader does not return.
  8- Subset-A holds a bounding box and a class for 125880 objects of 3622 frames. The classes are
  9  'S' for a sperm and 'Impurity' for everything else, at a ratio of about 27 to 1. Subset-C cuts one
 10  small image per box out of the frames, for a classification task. Read 'Subset-A/<frame>.xml' and
 11  'Subset-C/V<video>-F<frame>-<class><index>.png' if you need them.
 12- The masks of Subset-B carry a track id, because the file name of a mask ends with the id of the
 13  sperm and that id follows the same sperm over the frames of a video. This loader keeps the id as
 14  the label value, so the labels of one video track the sperm through time.
 15
 16NOTE: Only 10 of the 101 videos carry masks. The other 91 videos hold boxes alone.
 17
 18NOTE: The archive is a rar file, so the extraction needs the 'rarfile' package and the 'unrar'
 19program. The masks are stored as one full frame image per sperm, which is why the archive holds
 2025966 mask files for 451 frames.
 21
 22NOTE: The publication defines no split. This loader splits by video, because the frames of one video
 23show the same sperm over time and a split over frames would leak.
 24
 25The dataset is located at https://doi.org/10.6084/m9.figshare.15074253.v1.
 26This dataset is from the publication https://doi.org/10.1016/j.bbe.2021.12.010.
 27Please cite it if you use this dataset in your research. The authors welcome non-commercial research
 28work on this data.
 29"""
 30
 31import os
 32import re
 33from glob import glob
 34from pathlib import Path
 35from natsort import natsorted
 36from typing import List, Literal, Optional, Sequence, Tuple, Union
 37
 38import numpy as np
 39import imageio.v3 as imageio
 40
 41from torch.utils.data import DataLoader, Dataset
 42
 43import torch_em
 44
 45from .. import util
 46
 47
 48URL = "https://ndownloader.figshare.com/files/28986378"
 49CHECKSUM = "f6956e21eab440806f76a9805655b8a7e56f2555e8f6a03983d2eebab4b2ba71"
 50
 51ARCHIVE_ROOT = "Data Set"
 52
 53# The ten videos of Subset-B, and how many annotated frames each of them holds.
 54VIDEOS = {
 55    "S_0001": 45, "S_0003": 38, "S_0005": 36, "S_0006": 26, "S_0007": 10,
 56    "S_0008": 90, "S_0009": 48, "S_0010": 32, "S_0011": 39, "S_0012": 87,
 57}
 58
 59# A video belongs to one split only, so the frames of a video never spread over two splits.
 60SPLITS = {
 61    "train": ("S_0001", "S_0005", "S_0008", "S_0010", "S_0011"),  # 242 frames
 62    "val": ("S_0003", "S_0007"),  # 48 frames
 63    "test": ("S_0006", "S_0009", "S_0012"),  # 161 frames
 64}
 65
 66# The masks are not clean binary images, so the loader compares against the middle of the range.
 67MASK_THRESHOLD = 127
 68
 69
 70def _extract_archive(archive_path: str, path: str) -> None:
 71    """Extract the frames and the masks of Subset-B out of the rar archive."""
 72    try:
 73        import rarfile
 74    except ImportError:
 75        raise RuntimeError(
 76            "The 'rarfile' package is required to extract the SVIA archive. "
 77            "Install it with 'pip install rarfile', and install the 'unrar' program as well."
 78        )
 79
 80    with rarfile.RarFile(archive_path) as archive:
 81        members = [
 82            name for name in archive.namelist()
 83            if name.startswith(f"{ARCHIVE_ROOT}/Subset-B/")
 84            or name.startswith(f"{ARCHIVE_ROOT}/Frames from original videos/")
 85        ]
 86        if not members:
 87            raise RuntimeError(f"The archive {archive_path} does not hold Subset-B.")
 88        archive.extractall(path, members=members)
 89
 90
 91def _get_track_id(mask_stem: str, frame_name: str) -> int:
 92    """Read the track id of a sperm out of the name of its mask.
 93
 94    A mask of the frame 'S_0001_0026' is normally named 'S_0001_0026_0010'. Two files of the
 95    archive break that rule, 'S_0001_0026-0010' and 'S_0003_001311png', so the id comes from the
 96    part of the name that follows the name of the frame.
 97    """
 98    remainder = mask_stem[len(frame_name):] if mask_stem.startswith(frame_name) else mask_stem
 99    digits = re.findall(r"\d+", remainder)
100    if not digits:
101        raise RuntimeError(f"Could not read a track id from the mask '{mask_stem}'.")
102    return int(digits[0])
103
104
105def _create_instance_labels(data_dir: str, video: str) -> str:
106    """Merge the per sperm masks of every frame into one instance label image."""
107    from tqdm import tqdm
108
109    label_dir = os.path.join(data_dir, "instance_labels", video)
110    os.makedirs(label_dir, exist_ok=True)
111
112    frame_dirs = natsorted(glob(os.path.join(data_dir, "Subset-B", video, "*")))
113    for frame_dir in tqdm(frame_dirs, desc=f"Preprocess '{video}'"):
114        name = os.path.basename(frame_dir)
115        output_path = os.path.join(label_dir, f"{name}.tif")
116        if os.path.exists(output_path):
117            continue
118
119        image_path = os.path.join(data_dir, "Frames from original videos", video, f"{name}.png")
120        if not os.path.exists(image_path):
121            continue
122
123        shape = imageio.imread(image_path).shape[:2]
124        labels = np.zeros(shape, dtype="uint16")
125        for mask_path in natsorted(glob(os.path.join(frame_dir, "*.png"))):
126            track_id = _get_track_id(Path(mask_path).stem, name)
127            mask = imageio.imread(mask_path)
128            if mask.ndim == 3:
129                mask = mask[..., 0]
130            labels[mask > MASK_THRESHOLD] = track_id
131
132        imageio.imwrite(output_path, labels, compression="zlib")
133
134    return label_dir
135
136
137def get_svia_data(path: Union[os.PathLike, str], download: bool = False) -> str:
138    """Download the SVIA dataset.
139
140    Args:
141        path: Filepath to a folder where the downloaded data will be saved.
142        download: Whether to download the data if it is not present.
143
144    Returns:
145        The filepath to the extracted data.
146    """
147    data_dir = os.path.join(path, ARCHIVE_ROOT)
148    if os.path.exists(data_dir):
149        return data_dir
150
151    os.makedirs(path, exist_ok=True)
152    archive_path = os.path.join(path, "svia.rar")
153    util.download_source(archive_path, URL, download, CHECKSUM)
154    _extract_archive(archive_path, path)
155
156    return data_dir
157
158
159def get_svia_paths(
160    path: Union[os.PathLike, str],
161    split: Optional[Literal["train", "val", "test"]] = "train",
162    videos: Optional[Sequence[str]] = None,
163    download: bool = False,
164) -> Tuple[List[str], List[str]]:
165    """Get paths to the SVIA data.
166
167    Args:
168        path: Filepath to a folder where the downloaded data will be saved.
169        split: The data split, which groups the videos. Ignored when you pass `videos`.
170        videos: The videos to use, for example ('S_0001',). Overrides `split`.
171        download: Whether to download the data if it is not present.
172
173    Returns:
174        List of filepaths for the image data.
175        List of filepaths for the label data.
176    """
177    if videos is None:
178        if split not in SPLITS:
179            raise ValueError(f"'{split}' is not a valid split. Choose from {list(SPLITS)}, or pass videos.")
180        videos = SPLITS[split]
181    else:
182        for video in videos:
183            if video not in VIDEOS:
184                raise ValueError(f"'{video}' has no masks. Choose from {list(VIDEOS)}.")
185
186    data_dir = get_svia_data(path, download)
187
188    image_paths, label_paths = [], []
189    for video in videos:
190        label_dir = _create_instance_labels(data_dir, video)
191        for label_path in natsorted(glob(os.path.join(label_dir, "*.tif"))):
192            name = Path(label_path).stem
193            image_path = os.path.join(data_dir, "Frames from original videos", video, f"{name}.png")
194            if not os.path.exists(image_path):
195                continue
196            image_paths.append(image_path)
197            label_paths.append(label_path)
198
199    if not image_paths:
200        raise RuntimeError(f"Could not find any SVIA data in {data_dir}.")
201
202    return image_paths, label_paths
203
204
205def get_svia_dataset(
206    path: Union[os.PathLike, str],
207    patch_shape: Tuple[int, int],
208    split: Optional[Literal["train", "val", "test"]] = "train",
209    videos: Optional[Sequence[str]] = None,
210    offsets: Optional[List[List[int]]] = None,
211    boundaries: bool = False,
212    binary: bool = False,
213    download: bool = False,
214    **kwargs,
215) -> Dataset:
216    """Get the SVIA dataset for sperm segmentation.
217
218    Args:
219        path: Filepath to a folder where the downloaded data will be saved.
220        patch_shape: The 2D patch shape to use for training.
221        split: The data split, which groups the videos. Ignored when you pass `videos`.
222        videos: The videos to use, for example ('S_0001',). Overrides `split`.
223        offsets: Offset values for affinity computation used as target.
224        boundaries: Whether to compute boundaries as the target.
225        binary: Whether to use a binary segmentation target.
226        download: Whether to download the data if it is not present.
227        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
228
229    Returns:
230        The segmentation dataset.
231    """
232    if len(patch_shape) != 2:
233        raise ValueError(f"The SVIA patch shape must be two-dimensional, got {patch_shape}.")
234
235    image_paths, label_paths = get_svia_paths(path, split, videos, download)
236
237    kwargs, _ = util.add_instance_label_transform(
238        kwargs, add_binary_target=True, offsets=offsets, boundaries=boundaries, binary=binary,
239    )
240    kwargs = util.ensure_transforms(ndim=2, **kwargs)
241
242    return torch_em.default_segmentation_dataset(
243        raw_paths=image_paths,
244        raw_key=None,
245        label_paths=label_paths,
246        label_key=None,
247        patch_shape=patch_shape,
248        is_seg_dataset=False,
249        ndim=2,
250        **kwargs,
251    )
252
253
254def get_svia_loader(
255    path: Union[os.PathLike, str],
256    batch_size: int,
257    patch_shape: Tuple[int, int],
258    split: Optional[Literal["train", "val", "test"]] = "train",
259    videos: Optional[Sequence[str]] = None,
260    offsets: Optional[List[List[int]]] = None,
261    boundaries: bool = False,
262    binary: bool = False,
263    download: bool = False,
264    **kwargs,
265) -> DataLoader:
266    """Get the SVIA dataloader for sperm segmentation.
267
268    Args:
269        path: Filepath to a folder where the downloaded data will be saved.
270        batch_size: The batch size for training.
271        patch_shape: The 2D patch shape to use for training.
272        split: The data split, which groups the videos. Ignored when you pass `videos`.
273        videos: The videos to use, for example ('S_0001',). Overrides `split`.
274        offsets: Offset values for affinity computation used as target.
275        boundaries: Whether to compute boundaries as the target.
276        binary: Whether to use a binary segmentation target.
277        download: Whether to download the data if it is not present.
278        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or the PyTorch DataLoader.
279
280    Returns:
281        The DataLoader.
282    """
283    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
284    dataset = get_svia_dataset(
285        path=path,
286        patch_shape=patch_shape,
287        split=split,
288        videos=videos,
289        offsets=offsets,
290        boundaries=boundaries,
291        binary=binary,
292        download=download,
293        **ds_kwargs,
294    )
295    return torch_em.get_data_loader(dataset, batch_size=batch_size, **loader_kwargs)