torch_em.data.datasets.medical.camus

  1import os
  2from glob import glob
  3from typing import Union, Tuple, Optional
  4
  5import torch_em
  6from torch_em.transform.generic import ResizeInputs
  7
  8from .. import util
  9
 10
 11URL = "https://humanheart-project.creatis.insa-lyon.fr/database/api/v1/folder/63fde55f73e9f004868fb7ac/download"
 12CHECKSUM = "43745d640db5d979332bda7f00f4746747a2591b46efc8f1966b573ce8d65655"
 13
 14
 15def get_camus_data(path, download):
 16    os.makedirs(path, exist_ok=True)
 17
 18    data_dir = os.path.join(path, "database_nifti")
 19    if os.path.exists(data_dir):
 20        return data_dir
 21
 22    zip_path = os.path.join(path, "CAMUS.zip")
 23    util.download_source(path=zip_path, url=URL, download=download, checksum=CHECKSUM)
 24    util.unzip(zip_path=zip_path, dst=path)
 25
 26    return data_dir
 27
 28
 29def _get_camus_paths(path, chamber, download):
 30    data_dir = get_camus_data(path=path, download=download)
 31
 32    if chamber is None:
 33        chamber = "*"  # 2CH / 4CH
 34    else:
 35        assert chamber in [2, 4], f"{chamber} is  not a valid chamber choice for the acquisitions."
 36        chamber = f"{chamber}CH"
 37
 38    image_paths = sorted(glob(os.path.join(data_dir, "patient*", f"patient*_{chamber}_half_sequence.nii.gz")))
 39    gt_paths = sorted(glob(os.path.join(data_dir, "patient*", f"patient*_{chamber}_half_sequence_gt.nii.gz")))
 40
 41    return image_paths, gt_paths
 42
 43
 44def get_camus_dataset(
 45    path: Union[os.PathLike, str],
 46    patch_shape: Tuple[int, int],
 47    chamber: Optional[int] = None,
 48    resize_inputs: bool = False,
 49    download: bool = False,
 50    **kwargs
 51):
 52    """Dataset for segmenting cardiac structures in 2d echocardiography images.
 53
 54    The database is located at:
 55    https://humanheart-project.creatis.insa-lyon.fr/database/#collection/6373703d73e9f0047faa1bc8
 56
 57    This dataset is from the CAMUS challenge - https://doi.org/10.1109/TMI.2019.2900516.
 58    Please cite it if you use this dataset for a publication.
 59    """
 60    image_paths, gt_paths = _get_camus_paths(path=path, chamber=chamber, download=download)
 61
 62    if resize_inputs:
 63        raw_trafo = ResizeInputs(target_shape=patch_shape, is_label=False)
 64        label_trafo = ResizeInputs(target_shape=patch_shape, is_label=True)
 65        patch_shape = None
 66    else:
 67        patch_shape = patch_shape
 68        raw_trafo, label_trafo = None, None
 69
 70    dataset = torch_em.default_segmentation_dataset(
 71        raw_paths=image_paths,
 72        raw_key="data",
 73        label_paths=gt_paths,
 74        label_key="data",
 75        patch_shape=patch_shape,
 76        raw_transform=raw_trafo,
 77        label_transform=label_trafo,
 78        **kwargs
 79    )
 80
 81    return dataset
 82
 83
 84def get_camus_loader(
 85    path: Union[os.PathLike, str],
 86    patch_shape: Tuple[int, int],
 87    batch_size: int,
 88    chamber: Optional[int] = None,
 89    resize_inputs: bool = False,
 90    download: bool = False,
 91    **kwargs
 92):
 93    """Dataloader for segmenting cardiac structures in 2d echocardiography images. See `get_camus_dataset` for details
 94    """
 95    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
 96    dataset = get_camus_dataset(
 97        path=path, patch_shape=patch_shape, chamber=chamber, resize_inputs=resize_inputs, download=download, **ds_kwargs
 98    )
 99    loader = torch_em.get_data_loader(dataset=dataset, batch_size=batch_size, **loader_kwargs)
100    return loader