torch_em.data.datasets.medical.fass

  1"""FASS is the Fetal Abdominal Structures Segmentation dataset for segmenting the
  2abdominal aorta artery, intrahepatic umbilical vein, stomach, and liver in fetal
  3abdominal circumference ultrasound images.
  4
  5The dataset is located at https://data.mendeley.com/datasets/4gcpm9dsc3/1 (CC BY 4.0).
  6This dataset is from the publication https://doi.org/10.17632/4gcpm9dsc3.1.
  7Please cite it if you use this dataset for your research.
  8
  9NOTE: Mendeley's "Download all files" link goes through a Cloudflare bot check that
 10blocks plain HTTP clients, so this module downloads the dataset's zip archive via its
 11`public-files` mirror url instead, which is not gated behind that check.
 12"""
 13
 14import os
 15from glob import glob
 16from tqdm import tqdm
 17from typing import Union, Tuple, List
 18
 19import numpy as np
 20import imageio.v3 as imageio
 21
 22from torch.utils.data import Dataset, DataLoader
 23
 24import torch_em
 25
 26from .. import util
 27
 28
 29URL = "https://data.mendeley.com/public-files/datasets/4gcpm9dsc3/files/89e74076-ff57-4e81-9634-4fc29c6128ff/file_downloaded"  # noqa
 30CHECKSUM = "bde4ab58689bdbc59fa8c0222d41a9fbe1edf16dab40ac95346b6c2f2a7d22ee"
 31
 32# The order in which the (mostly non-overlapping) per-structure binary masks are painted
 33# into the single-channel label. A structure painted later wins ties on the rare pixels
 34# where two structures' masks overlap.
 35STRUCTURES = ["liver", "stomach", "artery", "vein"]
 36
 37
 38def get_fass_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 39    """Download the FASS dataset.
 40
 41    Args:
 42        path: Filepath to a folder where the data is downloaded for further processing.
 43        download: Whether to download the data if it is not present.
 44
 45    Returns:
 46        Filepath to the folder with the downloaded images and structure annotations.
 47    """
 48    data_dir = os.path.join(path, "ARRAY_FORMAT")
 49    if os.path.exists(data_dir):
 50        return path
 51
 52    os.makedirs(path, exist_ok=True)
 53
 54    zip_path = os.path.join(path, "fass.zip")
 55    util.download_source(path=zip_path, url=URL, download=download, checksum=CHECKSUM)
 56    util.unzip(zip_path=zip_path, dst=path)
 57
 58    return path
 59
 60
 61def get_fass_paths(path: Union[os.PathLike, str], download: bool = False) -> Tuple[List[str], List[str]]:
 62    """Get paths to the FASS data.
 63
 64    Args:
 65        path: Filepath to a folder where the data is downloaded for further processing.
 66        download: Whether to download the data if it is not present.
 67
 68    Returns:
 69        List of filepaths for the image data.
 70        List of filepaths for the label data.
 71    """
 72    data_dir = get_fass_data(path=path, download=download)
 73
 74    image_paths = sorted(glob(os.path.join(data_dir, "IMAGES", "*.png")))
 75
 76    label_dir = os.path.join(data_dir, "labels")
 77    os.makedirs(label_dir, exist_ok=True)
 78
 79    gt_paths = []
 80    for image_path in tqdm(image_paths, desc="Preprocessing FASS labels"):
 81        fname = os.path.splitext(os.path.basename(image_path))[0]
 82        gt_path = os.path.join(label_dir, f"{fname}.tif")
 83        gt_paths.append(gt_path)
 84        if os.path.exists(gt_path):
 85            continue
 86
 87        npy_path = os.path.join(data_dir, "ARRAY_FORMAT", f"{fname}.npy")
 88        structures = np.load(npy_path, allow_pickle=True).item()["structures"]
 89
 90        label = np.zeros(structures["liver"].shape, dtype="uint8")
 91        for i, name in enumerate(STRUCTURES, start=1):
 92            label[structures[name] > 0] = i
 93
 94        imageio.imwrite(gt_path, label, compression="zlib")
 95
 96    return image_paths, gt_paths
 97
 98
 99def get_fass_dataset(
100    path: Union[os.PathLike, str],
101    patch_shape: Tuple[int, int],
102    resize_inputs: bool = False,
103    download: bool = False,
104    **kwargs
105) -> Dataset:
106    """Get the FASS dataset for fetal abdominal structure segmentation.
107
108    Args:
109        path: Filepath to a folder where the data is downloaded for further processing.
110        patch_shape: The patch shape to use for training.
111        resize_inputs: Whether to resize the inputs to the patch shape.
112        download: Whether to download the data if it is not present.
113        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
114
115    Returns:
116        The segmentation dataset.
117    """
118    image_paths, gt_paths = get_fass_paths(path, download)
119
120    if resize_inputs:
121        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": True}
122        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
123            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
124        )
125
126    return torch_em.default_segmentation_dataset(
127        raw_paths=image_paths,
128        raw_key=None,
129        label_paths=gt_paths,
130        label_key=None,
131        patch_shape=patch_shape,
132        is_seg_dataset=False,
133        **kwargs
134    )
135
136
137def get_fass_loader(
138    path: Union[os.PathLike, str],
139    patch_shape: Tuple[int, int],
140    batch_size: int,
141    resize_inputs: bool = False,
142    download: bool = False,
143    **kwargs
144) -> DataLoader:
145    """Get the FASS dataloader for fetal abdominal structure segmentation.
146
147    Args:
148        path: Filepath to a folder where the data is downloaded for further processing.
149        patch_shape: The patch shape to use for training.
150        batch_size: The batch size for training.
151        resize_inputs: Whether to resize the inputs to the patch shape.
152        download: Whether to download the data if it is not present.
153        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
154
155    Returns:
156        The DataLoader.
157    """
158    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
159    dataset = get_fass_dataset(path, patch_shape, resize_inputs, download, **ds_kwargs)
160    return torch_em.get_data_loader(dataset=dataset, batch_size=batch_size, **loader_kwargs)