torch_em.data.datasets.light_microscopy.fl2net

  1"""The FL2-Net dataset contains 3D bright-field microscopy timelapses of mouse embryos
  2with nuclear instance segmentation annotations.
  3
  4The dataset holds 84 embryos with 506 timepoints each, so 42504 volumes in total. The authors split
  5the data by embryo, and no embryo appears in more than one split. The images are label-free, which
  6makes the nuclei much harder to see than in a fluorescence image.
  7
  8NOTE: Every volume holds 51 z-slices, but each embryo was cropped to its own field of view, so the
  9xy shape differs between embryos. It ranges from (92, 102) to (158, 158), and is not square for most
 10embryos. Volumes smaller than `patch_shape` are padded, so keep the last two entries of your
 11`patch_shape` at 92 or below to train on unpadded data.
 12
 13NOTE: The volumes are stored as uncompressed uint16 tif files of about 1.8 MB each, so extracted in
 14full each of the two archives needs about 76 GB. The loader therefore extracts only the timepoints
 15that you request. Use `stride` to set how many timepoints it skips, or pass `timepoints` to select
 16them. The default stride of 25 keeps 21 of the 506 timepoints, which comes out at about 2 GB for the
 17train split.
 18
 19NOTE: The archives store their files in an arbitrary order, so extracting even a few timepoints
 20reads through the whole archive once. This takes about a minute for the annotations and much longer
 21for the images. The files are cached on disk, so this cost is only paid the first time.
 22
 23NOTE: The images take up 64 GiB as a single archive, so expect the download to run for a while.
 24Google Drive refuses that archive to anonymous callers with a 'quota exceeded' html page, which it
 25serves under HTTP 200 rather than as an error. It does answer ranged requests though, so this module
 26downloads the archives range by range, which also lets an interrupted download resume. If the
 27download fails anyway, fetch both archives manually from the links in
 28https://github.com/funalab/FL2-Net and place them in `path` as 'raw.tar.gz' and 'gt.tar.gz'.
 29
 30The dataset is located at https://github.com/funalab/FL2-Net.
 31This dataset is from the publication https://doi.org/10.1016/j.compbiomed.2025.111179.
 32Please cite it if you use this dataset in your research.
 33"""
 34
 35import os
 36import re
 37import tarfile
 38from tqdm import tqdm
 39from typing import List, Literal, Optional, Sequence, Tuple, Union
 40
 41import requests
 42
 43from torch.utils.data import DataLoader, Dataset
 44
 45import torch_em
 46
 47from .. import util
 48
 49
 50# These are the links from the dataset section of the FL2-Net README. Google Drive answers them with
 51# a 'quota exceeded' html page rather than the file, so the download below addresses the files by id
 52# instead. The urls are kept for reference and for the manual download instructions.
 53URLS = {
 54    "images": "https://drive.usercontent.google.com/download?id=1OAMmFM76TputGnU6nell6LU81N0hDmRc&confirm=xxx",
 55    "labels": "https://drive.usercontent.google.com/download?id=1hdSnCthLtyKMCahFLHUz36Awtj2-OC6T&confirm=xxx",
 56}
 57
 58FILE_IDS = {"images": "1OAMmFM76TputGnU6nell6LU81N0hDmRc", "labels": "1hdSnCthLtyKMCahFLHUz36Awtj2-OC6T"}
 59
 60CHECKSUMS = {
 61    "images": None,  # Filled in once the 64 GiB download has completed and been verified.
 62    "labels": "9c12b70978f3995662f377dac8fc173abdc0a350ee3c38e6367096c87c2d2200",
 63}
 64
 65# The size of the archives in bytes, as reported by Google Drive. The download checks against these,
 66# because a truncated transfer is otherwise only caught by the much slower checksum.
 67ARCHIVE_SIZES = {"images": 69032222254, "labels": 456938451}
 68
 69MANUAL_DOWNLOAD_MESSAGE = (
 70    "Google Drive refused to serve the FL2-Net archives, because too many users have downloaded them "
 71    "recently. Please download '{name}' manually from the dataset links in "
 72    "https://github.com/funalab/FL2-Net, save it as '{path}', and run this function again."
 73)
 74
 75ARCHIVE_NAMES = {"images": "raw.tar.gz", "labels": "gt.tar.gz"}
 76
 77N_TIMEPOINTS = 506
 78
 79# The authors split the data by embryo in datasets/split_list_411 of the FL2-Net repository.
 80SPLITS = {
 81    "train": (
 82        "F001/Embryo01", "F001/Embryo02", "F001/Embryo03", "F001/Embryo04",
 83        "F001/Embryo06", "F001/Embryo08", "F001/Embryo09", "F001/Embryo10",
 84        "F002/Embryo01", "F002/Embryo03", "F002/Embryo04", "F002/Embryo06",
 85        "F002/Embryo07", "F002/Embryo08", "F002/Embryo10", "F002/Embryo11",
 86        "F003/Embryo01", "F003/Embryo02", "F003/Embryo04", "F003/Embryo05",
 87        "F003/Embryo08", "F003/Embryo09", "F003/Embryo10", "F003/Embryo12",
 88        "F004/Embryo01", "F004/Embryo02", "F004/Embryo05", "F004/Embryo06",
 89        "F004/Embryo08", "F004/Embryo09", "F004/Embryo10", "F004/Embryo12",
 90        "F005/Embryo02", "F005/Embryo04", "F005/Embryo05", "F005/Embryo06",
 91        "F005/Embryo08", "F005/Embryo09", "F005/Embryo10", "F005/Embryo11",
 92        "F006/Embryo01", "F006/Embryo04", "F006/Embryo05", "F006/Embryo08",
 93        "F006/Embryo09", "F006/Embryo10", "F006/Embryo11", "F006/Embryo12",
 94        "F007/Embryo03", "F007/Embryo04", "F007/Embryo05", "F007/Embryo06",
 95        "F007/Embryo07", "F007/Embryo09", "F007/Embryo10", "F007/Embryo11",
 96    ),
 97    "val": (
 98        "F001/Embryo05", "F001/Embryo12", "F002/Embryo05", "F002/Embryo09",
 99        "F003/Embryo03", "F003/Embryo06", "F004/Embryo04", "F004/Embryo11",
100        "F005/Embryo01", "F005/Embryo03", "F006/Embryo06", "F006/Embryo07",
101        "F007/Embryo01", "F007/Embryo12",
102    ),
103    "test": (
104        "F001/Embryo07", "F001/Embryo11", "F002/Embryo02", "F002/Embryo12",
105        "F003/Embryo07", "F003/Embryo11", "F004/Embryo03", "F004/Embryo07",
106        "F005/Embryo07", "F005/Embryo12", "F006/Embryo02", "F006/Embryo03",
107        "F007/Embryo02", "F007/Embryo08",
108    ),
109}
110
111
112def _get_download_url(session: requests.Session, file_id: str) -> str:
113    """Get a download url for a large Google Drive file, and store the matching cookie in `session`.
114
115    Google cannot virus scan files of this size, so it answers with an interstitial page that holds a
116    confirmation token. That token, together with the cookie, is what makes the file downloadable.
117    """
118    response = session.get(f"https://drive.google.com/uc?export=download&id={file_id}", timeout=120)
119    response.raise_for_status()
120    token = re.search(r'name="uuid" value="([^"]+)"', response.text)
121    if token is None:
122        raise RuntimeError(
123            "Google Drive did not return a download token for the FL2-Net archive. "
124            f"It answered with: {response.text[:200]!r}"
125        )
126    return (
127        f"https://drive.usercontent.google.com/download?id={file_id}"
128        f"&export=download&confirm=t&uuid={token.group(1)}"
129    )
130
131
132def _download_from_gdrive(path: str, file_id: str, total: int, checksum: Optional[str], desc: str) -> None:
133    """Download a large public Google Drive file in chunks, and resume an interrupted download.
134
135    Google refuses these files to anonymous callers with a 'quota exceeded' html page under HTTP 200,
136    so a plain download writes that page to disk instead of the file. A ranged request for the same
137    url is served normally, so this reads the file range by range. That also makes the download
138    resumable, which matters for an archive of this size.
139
140    Ranges of more than 512 MiB are refused the same way as an unranged request, so the chunk size
141    stays well below that, and halves whenever a chunk is refused in case the limit is lowered.
142    """
143    chunk_size = 256 * 1024**2
144    min_chunk_size = 32 * 1024**2
145    tmp_path = f"{path}.incomplete"
146    session = requests.Session()
147    url = _get_download_url(session, file_id)
148
149    with tqdm(total=total, unit="B", unit_scale=True, desc=desc) as progress:
150        offset = os.path.getsize(tmp_path) if os.path.exists(tmp_path) else 0
151        progress.update(offset)
152
153        while offset < total:
154            end = min(offset + chunk_size, total) - 1
155            expected = end - offset + 1
156            response = session.get(url, headers={"Range": f"bytes={offset}-{end}"}, stream=True, timeout=3600)
157
158            # A refusal comes back as HTTP 200 with an html body rather than as an error code, so the
159            # status and the length are both checked before anything is written.
160            if response.status_code != 206 or int(response.headers.get("Content-Length", 0)) != expected:
161                response.close()
162                if chunk_size > min_chunk_size:
163                    chunk_size //= 2
164                url = _get_download_url(session, file_id)
165                continue
166
167            with open(tmp_path, "ab") as f:
168                for chunk in response.iter_content(chunk_size=1024**2):
169                    f.write(chunk)
170                    progress.update(len(chunk))
171            offset = os.path.getsize(tmp_path)
172
173    if offset != total:
174        raise RuntimeError(f"Downloaded {offset} bytes of {path}, but expected {total}.")
175
176    util._check_checksum(tmp_path, checksum)
177    os.replace(tmp_path, path)
178
179
180def _get_archive_root(archive_path: str) -> str:
181    """Read the name of the top level folder of an archive.
182
183    The two archives do not agree on this name: the images sit under 'raw' and the annotations under
184    'qcanet'. Neither is documented, so it is read here rather than assumed.
185    """
186    # 'r|gz' reads the archive as a stream. That is all this needs, and it avoids the seeks that
187    # 'r:gz' performs for every member, which are expensive on a gzip stream of this size.
188    with tarfile.open(archive_path, "r|gz") as archive:
189        for member in archive:
190            root = member.name.split("/")[0]
191            if root:
192                return root
193    raise RuntimeError(f"The archive {archive_path} is empty.")
194
195
196def _extract_members(archive_path: str, relative_names: Sequence[str], destination: str) -> None:
197    """Extract the given files from an archive in one pass, and drop the top level folder.
198
199    The archives store their files in an arbitrary order, so the whole archive has to be read to
200    find the requested ones. Files that were extracted before are skipped.
201    """
202    missing = {name for name in relative_names if not os.path.exists(os.path.join(destination, name))}
203    if not missing:
204        return
205
206    root = _get_archive_root(archive_path)
207    wanted = {f"{root}/{name}": name for name in missing}
208
209    found = set()
210    desc = f"Extract {len(missing)} files from {os.path.basename(archive_path)}"
211    with tarfile.open(archive_path, "r|gz") as archive:
212        with tqdm(total=len(wanted), desc=desc) as progress:
213            try:
214                for member in archive:
215                    target = wanted.get(member.name)
216                    if target is None:
217                        continue
218                    output_path = os.path.join(destination, target)
219                    os.makedirs(os.path.dirname(output_path), exist_ok=True)
220                    # Write to a temporary path first, so that an interrupted extraction is not
221                    # mistaken for a complete one when this function is called again.
222                    tmp_path = f"{output_path}.incomplete"
223                    with archive.extractfile(member) as source, open(tmp_path, "wb") as f:
224                        f.write(source.read())
225                    os.replace(tmp_path, output_path)
226                    found.add(target)
227                    progress.update(1)
228                    if len(found) == len(wanted):
229                        break
230            except (tarfile.ReadError, EOFError) as e:
231                # An incomplete download ends mid-stream. Everything up to that point was extracted,
232                # so say what happened rather than letting a bare gzip error surface.
233                raise RuntimeError(
234                    f"The archive {archive_path} ends before its end-of-stream marker, so the "
235                    f"download is incomplete. {len(found)} of {len(missing)} requested files were "
236                    f"found before it broke off. Delete the archive to download it again, or fetch "
237                    f"it manually from https://github.com/funalab/FL2-Net. The original error was: {e}"
238                ) from e
239
240    if found != missing:
241        raise RuntimeError(
242            f"Could not find {len(missing - found)} of {len(missing)} files in {archive_path}. "
243            f"The first missing file is '{sorted(missing - found)[0]}'."
244        )
245
246
247def _get_timepoints(timepoints: Optional[Sequence[int]], stride: int) -> List[int]:
248    """Resolve the requested timepoints. The timepoint index starts at one."""
249    if timepoints is not None:
250        selected = sorted({int(t) for t in timepoints})
251        if not selected:
252            raise ValueError("You have to request at least one timepoint.")
253        for timepoint in selected:
254            if not 1 <= timepoint <= N_TIMEPOINTS:
255                raise ValueError(f"The timepoint {timepoint} is outside the range 1 to {N_TIMEPOINTS}.")
256        return selected
257
258    if stride < 1:
259        raise ValueError(f"The stride must be at least one, got {stride}.")
260    return list(range(1, N_TIMEPOINTS + 1, stride))
261
262
263def get_fl2net_data(path: Union[os.PathLike, str], download: bool = False) -> str:
264    """Download the FL2-Net dataset.
265
266    NOTE: The image archive is 64 GiB, so this runs for a while. It can be interrupted and resumed.
267    Download the archives manually from the links in https://github.com/funalab/FL2-Net and place
268    them in `path` as 'raw.tar.gz' and 'gt.tar.gz' if the download fails.
269
270    Args:
271        path: Filepath to a folder where the downloaded data will be saved.
272        download: Whether to download the data if it is not present.
273
274    Returns:
275        The filepath to the folder that holds the archives.
276    """
277    os.makedirs(path, exist_ok=True)
278
279    for key, archive_name in ARCHIVE_NAMES.items():
280        archive_path = os.path.join(path, archive_name)
281        if os.path.exists(archive_path):
282            continue
283        if not download:
284            raise RuntimeError(f"Cannot find the data at {archive_path}, but download was set to False")
285
286        _download_from_gdrive(
287            path=archive_path,
288            file_id=FILE_IDS[key],
289            total=ARCHIVE_SIZES[key],
290            checksum=CHECKSUMS[key],
291            desc=f"Download {archive_name}",
292        )
293
294        # The download only ever writes ranges that Google served as file content, but a corrupt
295        # archive would otherwise not surface until the extraction fails with a confusing error.
296        if not tarfile.is_tarfile(archive_path):
297            os.remove(archive_path)
298            raise RuntimeError(MANUAL_DOWNLOAD_MESSAGE.format(name=archive_name, path=archive_path))
299
300    return path
301
302
303def get_fl2net_paths(
304    path: Union[os.PathLike, str],
305    split: Literal["train", "val", "test"] = "train",
306    embryos: Optional[Sequence[str]] = None,
307    timepoints: Optional[Sequence[int]] = None,
308    stride: int = 25,
309    download: bool = False,
310) -> Tuple[List[str], List[str]]:
311    """Get paths to the FL2-Net data.
312
313    Args:
314        path: Filepath to a folder where the downloaded data will be saved.
315        split: The data split. Either 'train', 'val' or 'test'.
316        embryos: The embryos to use, for example 'F001/Embryo01'. Defaults to all of the split.
317        timepoints: The timepoints to use, counted from one. Overrides `stride`.
318        stride: The number of timepoints to skip between two extractions.
319        download: Whether to download the data if it is not present.
320
321    Returns:
322        List of filepaths for the image data.
323        List of filepaths for the label data.
324    """
325    if split not in SPLITS:
326        raise ValueError(f"'{split}' is not a valid split. Choose from {list(SPLITS)}.")
327
328    if embryos is None:
329        embryos = SPLITS[split]
330    else:
331        for embryo in embryos:
332            if embryo not in SPLITS[split]:
333                raise ValueError(f"The embryo '{embryo}' is not part of the '{split}' split.")
334
335    get_fl2net_data(path, download)
336    selected = _get_timepoints(timepoints, stride)
337    # The names are sorted here, so that the images and the labels stay paired up below.
338    relative_names = [f"{embryo}/{timepoint:03d}.tif" for embryo in sorted(embryos) for timepoint in selected]
339
340    image_dir = os.path.join(path, "images")
341    label_dir = os.path.join(path, "labels")
342    _extract_members(os.path.join(path, ARCHIVE_NAMES["images"]), relative_names, image_dir)
343    _extract_members(os.path.join(path, ARCHIVE_NAMES["labels"]), relative_names, label_dir)
344
345    image_paths = [os.path.join(image_dir, name) for name in relative_names]
346    label_paths = [os.path.join(label_dir, name) for name in relative_names]
347    assert len(image_paths) == len(label_paths)
348    return image_paths, label_paths
349
350
351def get_fl2net_dataset(
352    path: Union[os.PathLike, str],
353    patch_shape: Tuple[int, int, int],
354    split: Literal["train", "val", "test"] = "train",
355    embryos: Optional[Sequence[str]] = None,
356    timepoints: Optional[Sequence[int]] = None,
357    stride: int = 25,
358    offsets: Optional[List[List[int]]] = None,
359    boundaries: bool = False,
360    binary: bool = False,
361    download: bool = False,
362    **kwargs,
363) -> Dataset:
364    """Get the FL2-Net dataset for nucleus segmentation.
365
366    Args:
367        path: Filepath to a folder where the downloaded data will be saved.
368        patch_shape: The 3D patch shape to use for training.
369        split: The data split. Either 'train', 'val' or 'test'.
370        embryos: The embryos to use, for example 'F001/Embryo01'. Defaults to all of the split.
371        timepoints: The timepoints to use, counted from one. Overrides `stride`.
372        stride: The number of timepoints to skip between two extractions.
373        offsets: Offset values for affinity computation used as target.
374        boundaries: Whether to compute boundaries as the target.
375        binary: Whether to use a binary segmentation target.
376        download: Whether to download the data if it is not present.
377        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
378
379    Returns:
380        The segmentation dataset.
381    """
382    if len(patch_shape) != 3:
383        raise ValueError(f"The FL2-Net patch shape must be three-dimensional, got {patch_shape}.")
384
385    image_paths, label_paths = get_fl2net_paths(path, split, embryos, timepoints, stride, download)
386
387    kwargs, _ = util.add_instance_label_transform(
388        kwargs, add_binary_target=True, offsets=offsets, boundaries=boundaries, binary=binary,
389    )
390    kwargs = util.ensure_transforms(ndim=3, **kwargs)
391
392    return torch_em.default_segmentation_dataset(
393        raw_paths=image_paths,
394        raw_key=None,
395        label_paths=label_paths,
396        label_key=None,
397        patch_shape=patch_shape,
398        ndim=3,
399        **kwargs,
400    )
401
402
403def get_fl2net_loader(
404    path: Union[os.PathLike, str],
405    batch_size: int,
406    patch_shape: Tuple[int, int, int],
407    split: Literal["train", "val", "test"] = "train",
408    embryos: Optional[Sequence[str]] = None,
409    timepoints: Optional[Sequence[int]] = None,
410    stride: int = 25,
411    offsets: Optional[List[List[int]]] = None,
412    boundaries: bool = False,
413    binary: bool = False,
414    download: bool = False,
415    **kwargs,
416) -> DataLoader:
417    """Get the FL2-Net dataloader for nucleus segmentation.
418
419    Args:
420        path: Filepath to a folder where the downloaded data will be saved.
421        batch_size: The batch size for training.
422        patch_shape: The 3D patch shape to use for training.
423        split: The data split. Either 'train', 'val' or 'test'.
424        embryos: The embryos to use, for example 'F001/Embryo01'. Defaults to all of the split.
425        timepoints: The timepoints to use, counted from one. Overrides `stride`.
426        stride: The number of timepoints to skip between two extractions.
427        offsets: Offset values for affinity computation used as target.
428        boundaries: Whether to compute boundaries as the target.
429        binary: Whether to use a binary segmentation target.
430        download: Whether to download the data if it is not present.
431        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or the PyTorch DataLoader.
432
433    Returns:
434        The DataLoader.
435    """
436    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
437    dataset = get_fl2net_dataset(
438        path=path,
439        patch_shape=patch_shape,
440        split=split,
441        embryos=embryos,
442        timepoints=timepoints,
443        stride=stride,
444        offsets=offsets,
445        boundaries=boundaries,
446        binary=binary,
447        download=download,
448        **ds_kwargs,
449    )
450    return torch_em.get_data_loader(dataset, batch_size=batch_size, **loader_kwargs)