torch_em.data.datasets.medical.feta24

  1"""The FETA24 contains annotations for fetal brain tissue segmentation in MRI.
  2
  3This dataset is from FeTa 2024 Challenge: https://doi.org/10.5281/zenodo.11192452.
  4The dataset is from the publication https://doi.org/10.1038/s41597-021-00946-3.
  5Please cite it if you use this dataset in your publication.
  6"""
  7
  8import os
  9from glob import glob
 10from natsort import natsorted
 11from typing import Union, Tuple, List
 12
 13from torch.utils.data import Dataset, DataLoader
 14
 15import torch_em
 16
 17from .. import util
 18
 19
 20def get_feta24_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 21    """Obtain the FETA24 data.
 22    
 23    This function describes the download functionality and ensures your data has been downloaded in expected format.
 24
 25    The dataset is from the FeTa Challenge 2024 - https://fetachallenge.github.io/ (Task 1: Segmentation).
 26    A detailed description of the dataset is provided here: https://fetachallenge.github.io/pages/Data_description.
 27    To download the dataset, please follow the below mentioned steps:
 28    - Go to the section `1. Request access and download the FeTa 2024 data from the University Children's Hospital
 29    Zurich` at `https://fetachallenge.github.io/pages/Data_download`, which explains the steps to be a registered user
 30    in Synapse platform and expects the user to agree with the mentioned conditions.
 31    - While registration, the users are expected to provide some information
 32    (see https://fetachallenge.github.io/pages/Data_download for details).
 33    - Next, you can proceed with requesting access (by following provided instructions) at
 34    https://www.synapse.org/#!Synapse:syn25649159/wiki/610007.
 35
 36    Once you have access to the dataset, you can use the synapse client or the platform download option to get
 37    the zipped files. It contains 80 scans paired with their segmentations (more details in the challenge website).
 38
 39    Finally, you should provide the path to the parent directory where the zipfile is stored.
 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 where the data is downloaded.
 47    """
 48    if download:
 49        print("Download is not supported due to the challenge's setup. See 'get_feta24_data' for details.")
 50
 51    data_dir = os.path.join(path, "feta_2.3")
 52    if os.path.exists(data_dir):
 53        return data_dir
 54
 55    zip_path = os.path.join(path, "feta_2.3.zip")
 56    if not os.path.exists(zip_path):
 57        raise FileNotFoundError(f"The downloaded zip file was not found. Please download it and place it at '{path}'.")
 58
 59    util.unzip(zip_path=zip_path, dst=path)
 60
 61    return data_dir
 62
 63
 64def get_feta24_paths(path: Union[os.PathLike, str], download: bool = False) -> Tuple[List[str], List[str]]:
 65    """Get the FETA24 dataset.
 66
 67    Args:
 68        path: Filepath to a folder where the data is downloaded for further processing.
 69        download: Whether to download the data if it is not present.
 70
 71    Returns:
 72        List of filepaths for the image data.
 73        List of filepaths for the label data.
 74    """
 75    data_dir = get_feta24_data(path=path, download=download)
 76
 77    base_dir = os.path.join(data_dir, "sub-*", "anat")
 78    image_paths = natsorted(glob(os.path.join(base_dir, "sub-*_rec-*_T2w.nii.gz")))
 79    gt_paths = natsorted(glob(os.path.join(base_dir, "sub-*_rec-*_dseg.nii.gz")))
 80
 81    return image_paths, gt_paths
 82
 83
 84def get_feta24_dataset(
 85    path: Union[os.PathLike, str],
 86    patch_shape: Tuple[int, ...],
 87    resize_inputs: bool = False,
 88    download: bool = False,
 89    **kwargs
 90) -> Dataset:
 91    """Get the FETA24 dataset for fetal brain tissue segmentation.
 92
 93    Args:
 94        path: Filepath to a folder where the data is downloaded for further processing.
 95        patch_shape: The patch shape to use for training.
 96        resize_inputs: Whether to resize inputs to the desired patch shape.
 97        download: Whether to download the data if it is not present.
 98        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
 99
100    Returns:
101        The segmentation dataset.
102    """
103    image_paths, gt_paths = get_feta24_paths(path, download)
104
105    if resize_inputs:
106        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
107        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
108            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
109        )
110
111    return torch_em.default_segmentation_dataset(
112        raw_paths=image_paths,
113        raw_key="data",
114        label_paths=gt_paths,
115        label_key="data",
116        patch_shape=patch_shape,
117        **kwargs
118    )
119
120
121def get_feta24_loader(
122    path: Union[os.PathLike, str],
123    batch_size: int,
124    patch_shape: Tuple[int, ...],
125    resize_inputs: bool = False,
126    download: bool = False,
127    **kwargs
128) -> DataLoader:
129    """Get the FETA24 dataloader for fetal brain tissue segmentation.
130
131    Args:
132        path: Filepath to a folder where the data is downloaded for further processing.
133        batch_size: The batch size for training.
134        patch_shape: The patch shape to use for training.
135        resize_inputs: Whether to resize inputs to the desired patch shape.
136        download: Whether to download the data if it is not present.
137        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
138
139    Returns:
140        The DataLoader.
141    """
142    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
143    dataset = get_feta24_dataset(path, patch_shape, resize_inputs, download, **ds_kwargs)
144    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)