torch_em.data.datasets.medical.camus

  1"""The CAMUS dataset contains annotations for cardiac structures segmentation in 2d echocardiography images.
  2
  3The database is located at:
  4https://humanheart-project.creatis.insa-lyon.fr/database/#collection/6373703d73e9f0047faa1bc8.
  5This dataset is from the publication https://doi.org/10.1109/TMI.2019.2900516.
  6Please cite it if you use this dataset for a publication.
  7"""
  8
  9import os
 10from glob import glob
 11from typing import Union, Tuple, Optional, Literal, List
 12
 13from torch.utils.data import Dataset, DataLoader
 14
 15import torch_em
 16
 17from .. import util
 18
 19
 20URL = "https://humanheart-project.creatis.insa-lyon.fr/database/api/v1/folder/63fde55f73e9f004868fb7ac/download"
 21
 22# TODO: the checksums are different with each download, not sure why
 23# CHECKSUM = "43745d640db5d979332bda7f00f4746747a2591b46efc8f1966b573ce8d65655"
 24
 25
 26def get_camus_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 27    """Get the CAMUS dataset.
 28
 29    Args:
 30        path: Filepath to a folder where the data is downloaded for further processing.
 31        download: Whether to download the data if it is not present.
 32
 33    Returns:
 34        Filepath where the data is downloaded.
 35    """
 36    data_dir = os.path.join(path, "database_nifti")
 37    if os.path.exists(data_dir):
 38        return data_dir
 39
 40    os.makedirs(path, exist_ok=True)
 41
 42    zip_path = os.path.join(path, "CAMUS.zip")
 43    util.download_source(path=zip_path, url=URL, download=download, checksum=None)
 44    util.unzip(zip_path=zip_path, dst=path)
 45
 46    return data_dir
 47
 48
 49def get_camus_paths(
 50    path: Union[os.PathLike, str], chamber: Optional[Literal[2, 4]] = None, download: bool = False
 51) -> Tuple[List[str], List[str]]:
 52    """Get paths to the CAMUS data.
 53
 54    Args:
 55        path: Filepath to a folder where the data is downloaded for further processing.
 56        chamber: The choice of chamber.
 57        download: Whether to download the data if it is not present.
 58
 59    Returns:
 60        List of filepaths for the image data.
 61        List of filepaths for the label data.
 62    """
 63    data_dir = get_camus_data(path=path, download=download)
 64
 65    if chamber is None:
 66        chamber = "*"  # 2CH / 4CH
 67    else:
 68        assert chamber in [2, 4], f"{chamber} is  not a valid chamber choice for the acquisitions."
 69        chamber = f"{chamber}CH"
 70
 71    image_paths = sorted(glob(os.path.join(data_dir, "patient*", f"patient*_{chamber}_half_sequence.nii.gz")))
 72    gt_paths = sorted(glob(os.path.join(data_dir, "patient*", f"patient*_{chamber}_half_sequence_gt.nii.gz")))
 73
 74    return image_paths, gt_paths
 75
 76
 77def get_camus_dataset(
 78    path: Union[os.PathLike, str],
 79    patch_shape: Tuple[int, int],
 80    chamber: Optional[Literal[2, 4]] = None,
 81    resize_inputs: bool = False,
 82    download: bool = False,
 83    **kwargs
 84) -> Dataset:
 85    """Get the CAMUS dataset for cardiac structure segmentation.
 86
 87    Args:
 88        path: Filepath to a folder where the data is downloaded for further processing.
 89        patch_shape: The patch shape to use for training.
 90        chamber: The choice of chamber.
 91        resize_inputs: Whether to resize inputs to the desired patch shape.
 92        download: Whether to download the data if it is not present.
 93        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
 94
 95    Returns:
 96        The segmentation dataset.
 97    """
 98    image_paths, gt_paths = get_camus_paths(path, chamber, download)
 99
100    if resize_inputs:
101        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
102        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
103            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
104        )
105
106    return torch_em.default_segmentation_dataset(
107        raw_paths=image_paths,
108        raw_key="data",
109        label_paths=gt_paths,
110        label_key="data",
111        patch_shape=patch_shape,
112        **kwargs
113    )
114
115
116def get_camus_loader(
117    path: Union[os.PathLike, str],
118    batch_size: int,
119    patch_shape: Tuple[int, int],
120    chamber: Optional[Literal[2, 4]] = None,
121    resize_inputs: bool = False,
122    download: bool = False,
123    **kwargs
124) -> DataLoader:
125    """Get the CAMUS dataloader for cardiac structure segmentation.
126
127    Args:
128        path: Filepath to a folder where the data is downloaded for further processing.
129        batch_size: The batch size for training.
130        patch_shape: The patch shape to use for training.
131        chamber: The choice of chamber.
132        resize_inputs: Whether to resize inputs to the desired patch shape.
133        download: Whether to download the data if it is not present.
134        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
135
136    Returns:
137        The DataLoader.
138    """
139    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
140    dataset = get_camus_dataset(path, patch_shape, chamber, resize_inputs, download, **ds_kwargs)
141    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)