torch_em.data.datasets.medical.lumase

  1"""The LumASe dataset contains annotations for lumbar vertebra anatomical substructure segmentation
  2in computed tomography (CT).
  3
  4The dataset consists of 663 individual vertebrae (L1 to L5) cropped from lumbar spine CT scans,
  5acquired at ShengJing Hospital of China Medical University with three different CT manufacturers
  6(Philips, Siemens and Toshiba). Cases with vertebral fractures, metallic implants, bone tumors or
  7other foreign materials are excluded. Each vertebra is voxel-wise annotated for 7 anatomical
  8substructures: superior articular process (SAP), vertebral body (VB), transverse process (TP),
  9lamina (L), pedicle (P), spinous process (SP) and inferior articular process (IAP).
 10
 11This dataset is located at https://doi.org/10.5281/zenodo.7181338.
 12This dataset is from the publication https://doi.org/10.1109/ISBI53787.2023.10230438.
 13The dataset is licensed under CC-BY-4.0.
 14Please cite the publication above if you use this dataset for your research.
 15"""
 16
 17import os
 18from glob import glob
 19from natsort import natsorted
 20from typing import Union, Tuple, List
 21
 22from torch.utils.data import Dataset, DataLoader
 23
 24import torch_em
 25
 26from .. import util
 27
 28
 29URL = "https://zenodo.org/records/7181338/files/L1-L5FineSegMix-663case.zip"
 30CHECKSUM = "61a280b446e1f1dd10935e0bed5cc8fa80152fc568f4d6d7dd28d7b857a322d0"
 31
 32
 33def get_lumase_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 34    """Download the LumASe dataset.
 35
 36    Args:
 37        path: Filepath to a folder where the data is downloaded for further processing.
 38        download: Whether to download the data if it is not present.
 39
 40    Returns:
 41        Filepath where the data is downloaded.
 42    """
 43    data_dir = os.path.join(path, "L1-L5FineSegMix-663case")
 44    if os.path.exists(data_dir):
 45        return data_dir
 46
 47    os.makedirs(path, exist_ok=True)
 48
 49    zip_path = os.path.join(path, "L1-L5FineSegMix-663case.zip")
 50    util.download_source(path=zip_path, url=URL, download=download, checksum=CHECKSUM)
 51    util.unzip(zip_path=zip_path, dst=path)
 52
 53    return data_dir
 54
 55
 56def get_lumase_paths(path: Union[os.PathLike, str], download: bool = False) -> Tuple[List[str], List[str]]:
 57    """Get paths to the LumASe data.
 58
 59    Args:
 60        path: Filepath to a folder where the data is downloaded for further processing.
 61        download: Whether to download the data if it is not present.
 62
 63    Returns:
 64        List of filepaths for the image data.
 65        List of filepaths for the label data.
 66    """
 67    data_dir = get_lumase_data(path, download)
 68
 69    label_paths = natsorted(glob(os.path.join(data_dir, "*_seg.nii.gz")))
 70    raw_paths = [p.replace("_seg.nii.gz", ".nii.gz") for p in label_paths]
 71    assert len(raw_paths) > 0 and all(os.path.exists(p) for p in raw_paths)
 72
 73    return raw_paths, label_paths
 74
 75
 76def get_lumase_dataset(
 77    path: Union[os.PathLike, str],
 78    patch_shape: Tuple[int, ...],
 79    resize_inputs: bool = False,
 80    download: bool = False,
 81    **kwargs
 82) -> Dataset:
 83    """Get the LumASe dataset for lumbar vertebra anatomical substructure segmentation.
 84
 85    Args:
 86        path: Filepath to a folder where the data is downloaded for further processing.
 87        patch_shape: The patch shape to use for training.
 88        resize_inputs: Whether to resize inputs to the desired patch shape.
 89        download: Whether to download the data if it is not present.
 90        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
 91
 92    Returns:
 93        The segmentation dataset.
 94    """
 95    raw_paths, label_paths = get_lumase_paths(path, download)
 96
 97    if resize_inputs:
 98        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
 99        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
100            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
101        )
102
103    return torch_em.default_segmentation_dataset(
104        raw_paths=raw_paths,
105        raw_key="data",
106        label_paths=label_paths,
107        label_key="data",
108        patch_shape=patch_shape,
109        is_seg_dataset=True,
110        **kwargs
111    )
112
113
114def get_lumase_loader(
115    path: Union[os.PathLike, str],
116    batch_size: int,
117    patch_shape: Tuple[int, ...],
118    resize_inputs: bool = False,
119    download: bool = False,
120    **kwargs
121) -> DataLoader:
122    """Get the LumASe dataloader for lumbar vertebra anatomical substructure segmentation.
123
124    Args:
125        path: Filepath to a folder where the data is downloaded for further processing.
126        batch_size: The batch size for training.
127        patch_shape: The patch shape to use for training.
128        resize_inputs: Whether to resize inputs to the desired patch shape.
129        download: Whether to download the data if it is not present.
130        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
131
132    Returns:
133        The DataLoader.
134    """
135    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
136    dataset = get_lumase_dataset(path, patch_shape, resize_inputs, download, **ds_kwargs)
137    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)