torch_em.data.datasets.medical.hc18

  1"""HC18 is a dataset for segmentation of the fetal head in 2d ultrasound images, together with
  2annotations for measuring the head circumference (HC).
  3
  4The dataset is located at https://doi.org/10.5281/zenodo.1327317, released under a CC BY 4.0 license.
  5This dataset is from the publication https://doi.org/10.1371/journal.pone.0200412. Please cite it if
  6you use this dataset in your research.
  7
  8NOTE: The dataset ships pre-rendered annotations as the ellipse fit of the head circumference drawn as
  9a thin closed contour (not a filled mask). We rasterize a filled binary mask from this contour by
 10filling the enclosed region (analogous in spirit to polygon rasterization, see `rasterize_rtstruct`
 11in `torch_em/data/datasets/util.py`, although here the shape is defined by a closed contour rather
 12than by polygon vertices).
 13
 14NOTE: Only the training split ships with ground truth annotations. The test split (available at
 15the same Zenodo record) only provides the raw images and pixel size metadata, without annotations,
 16and is therefore not supported by this loader.
 17"""
 18
 19import os
 20from glob import glob
 21from tqdm import tqdm
 22from natsort import natsorted
 23from typing import Union, Tuple, List
 24
 25import imageio.v3 as imageio
 26from scipy.ndimage import binary_fill_holes
 27
 28from torch.utils.data import Dataset, DataLoader
 29
 30import torch_em
 31
 32from .. import util
 33
 34
 35URL = "https://zenodo.org/records/1327317/files/training_set.zip"
 36CHECKSUM = "fd20d7909df892cfbdc0850de18072dbdad4dc3bc0a57202d4cc818d4715de36"
 37
 38
 39def get_hc18_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 40    """Download the HC18 dataset.
 41
 42    Args:
 43        path: Filepath to a folder where the data is downloaded for further processing.
 44        download: Whether to download the data if it is not present.
 45
 46    Returns:
 47        Filepath where the data is downloaded.
 48    """
 49    data_dir = os.path.join(path, "training_set")
 50    if os.path.exists(data_dir):
 51        return data_dir
 52
 53    os.makedirs(path, exist_ok=True)
 54
 55    zip_path = os.path.join(path, "training_set.zip")
 56    util.download_source(path=zip_path, url=URL, download=download, checksum=CHECKSUM)
 57    util.unzip(zip_path=zip_path, dst=path)
 58
 59    return data_dir
 60
 61
 62def _create_mask(annotation_path, gt_path):
 63    contour = imageio.imread(annotation_path)
 64    mask = binary_fill_holes(contour > 0).astype("uint8")
 65    imageio.imwrite(gt_path, mask)
 66
 67
 68def get_hc18_paths(path: Union[os.PathLike, str], download: bool = False) -> Tuple[List[str], List[str]]:
 69    """Get paths to the HC18 data.
 70
 71    Args:
 72        path: Filepath to a folder where the data is downloaded for further processing.
 73        download: Whether to download the data if it is not present.
 74
 75    Returns:
 76        List of filepaths for the image data.
 77        List of filepaths for the label data.
 78    """
 79    data_dir = get_hc18_data(path=path, download=download)
 80
 81    # NOTE: Some patients have multiple annotated images (eg. '010_HC.png' and '010_2HC.png'), so we
 82    # match on the 'HC.png' suffix rather than '_HC.png' and explicitly exclude the annotation images.
 83    image_paths = natsorted(
 84        p for p in glob(os.path.join(data_dir, "*HC.png")) if "Annotation" not in os.path.basename(p)
 85    )
 86
 87    gt_dir = os.path.join(path, "masks")
 88    os.makedirs(gt_dir, exist_ok=True)
 89
 90    gt_paths = []
 91    for image_path in tqdm(image_paths, desc="Rasterizing head circumference masks"):
 92        image_id = os.path.basename(image_path)[:-len(".png")]
 93        annotation_path = os.path.join(data_dir, f"{image_id}_Annotation.png")
 94        assert os.path.exists(annotation_path), f"The annotation for '{image_path}' is missing."
 95
 96        gt_path = os.path.join(gt_dir, f"{image_id}.tif")
 97        if not os.path.exists(gt_path):
 98            _create_mask(annotation_path, gt_path)
 99
100        gt_paths.append(gt_path)
101
102    assert len(image_paths) == len(gt_paths) == 999, len(image_paths)
103
104    return image_paths, gt_paths
105
106
107def get_hc18_dataset(
108    path: Union[os.PathLike, str],
109    patch_shape: Tuple[int, int],
110    resize_inputs: bool = False,
111    download: bool = False,
112    **kwargs
113) -> Dataset:
114    """Get the HC18 dataset for segmentation of the fetal head in ultrasound images.
115
116    Args:
117        path: Filepath to a folder where the data is downloaded for further processing.
118        patch_shape: The patch shape to use for training.
119        resize_inputs: Whether to resize the inputs to the patch shape.
120        download: Whether to download the data if it is not present.
121        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
122
123    Returns:
124        The segmentation dataset.
125    """
126    image_paths, gt_paths = get_hc18_paths(path, download)
127
128    if resize_inputs:
129        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
130        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
131            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
132        )
133
134    return torch_em.default_segmentation_dataset(
135        raw_paths=image_paths,
136        raw_key=None,
137        label_paths=gt_paths,
138        label_key=None,
139        patch_shape=patch_shape,
140        is_seg_dataset=False,
141        **kwargs
142    )
143
144
145def get_hc18_loader(
146    path: Union[os.PathLike, str],
147    batch_size: int,
148    patch_shape: Tuple[int, int],
149    resize_inputs: bool = False,
150    download: bool = False,
151    **kwargs
152) -> DataLoader:
153    """Get the HC18 dataloader for segmentation of the fetal head in ultrasound images.
154
155    Args:
156        path: Filepath to a folder where the data is downloaded for further processing.
157        batch_size: The batch size for training.
158        patch_shape: The patch shape to use for training.
159        resize_inputs: Whether to resize the inputs to the patch shape.
160        download: Whether to download the data if it is not present.
161        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
162
163    Returns:
164        The DataLoader.
165    """
166    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
167    dataset = get_hc18_dataset(path, patch_shape, resize_inputs, download, **ds_kwargs)
168    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)