torch_em.data.datasets.medical.jnuifm

  1import os
  2from glob import glob
  3from natsort import natsorted
  4from urllib.parse import urljoin
  5from typing import Union, Tuple, List
  6
  7from torch.utils.data import Dataset, DataLoader
  8
  9import torch_em
 10
 11from .. import util
 12
 13
 14BASE_URL = "https://zenodo.org/records/7851339/files/"
 15URL = urljoin(BASE_URL, "Pubic%20Symphysis-Fetal%20Head%20Segmentation%20and%20Angle%20of%20Progression.zip")
 16CHECKSUM = "2b14d1c78e11cfb799d74951b0b985b90777c195f7a456ccd00528bf02802e21"
 17
 18
 19def get_jnuifm_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 20    """Download the JNUIFM dataset.
 21
 22    Args:
 23        path: Filepath to a folder where the data is downloaded for further processing.
 24        download: Whether to download the data if it is not present.
 25
 26    Returns:
 27        Filepath where the data is downloaded.
 28    """
 29    data_dir = os.path.join(path, r"Pubic Symphysis-Fetal Head Segmentation and Angle of Progression")
 30    if os.path.exists(data_dir):
 31        return data_dir
 32
 33    os.makedirs(path, exist_ok=True)
 34
 35    zip_path = os.path.join(path, "JNU-IFM.zip")
 36    util.download_source(path=zip_path, url=URL, download=download, checksum=CHECKSUM)
 37    util.unzip(zip_path=zip_path, dst=path)
 38
 39    return data_dir
 40
 41
 42def get_jnuifm_paths(path: Union[os.PathLike, str], download: bool = False) -> Tuple[List[str], List[str]]:
 43    """Get paths to the JNUIFM data.
 44
 45    Args:
 46        path: Filepath to a folder where the data is downloaded for further processing.
 47        download: Whether to download the data if it is not present.
 48
 49    Returns:
 50        List of filepaths for the image data.
 51        List of filepaths for the label data.
 52    """
 53    data_dir = get_jnuifm_data(path, download)
 54    image_paths = natsorted(glob(os.path.join(data_dir, "image_mha", "*.mha")))
 55    gt_paths = natsorted(glob(os.path.join(data_dir, "label_mha", "*.mha")))
 56    return image_paths, gt_paths
 57
 58
 59def get_jnuifm_dataset(
 60    path: Union[os.PathLike, str],
 61    patch_shape: Tuple[int, int],
 62    resize_inputs: bool = False,
 63    download: bool = False,
 64    **kwargs
 65) -> Dataset:
 66    """Get the JNUIFM dataset for segmentation of pubic symphysis and fetal head in ultrasound images.
 67
 68    Args:
 69        path: Filepath to a folder where the data is downloaded for further processing.
 70        patch_shape: The patch shape to use for training.
 71        resize_inputs: Whether to resize the inputs to the expected patch shape.
 72        download: Whether to download the data if it is not present.
 73        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
 74
 75    Returns:
 76        The segmentation dataset.
 77    """
 78    image_paths, gt_paths = get_jnuifm_paths(path, download)
 79
 80    if resize_inputs:
 81        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": True}
 82        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
 83            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
 84        )
 85
 86    return torch_em.default_segmentation_dataset(
 87        raw_paths=image_paths,
 88        raw_key=None,
 89        label_paths=gt_paths,
 90        label_key=None,
 91        patch_shape=patch_shape,
 92        ndim=2,
 93        with_channels=True,
 94        is_seg_dataset=False,
 95        **kwargs
 96    )
 97
 98
 99def get_jnuifm_loader(
100    path: Union[os.PathLike, str],
101    batch_size: int,
102    patch_shape: Tuple[int, int],
103    resize_inputs: bool = False,
104    download: bool = False,
105    **kwargs
106) -> DataLoader:
107    """Get the JNUIFM dataloader for segmentation of pubic symphysis and fetal head in ultrasound images.
108
109    Args:
110        path: Filepath to a folder where the data is downloaded for further processing.
111        batch_size: The batch size for training.
112        patch_shape: The patch shape to use for training.
113        resize_inputs: Whether to resize the inputs to the expected patch shape.
114        download: Whether to download the data if it is not present.
115        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
116
117    Returns:
118        The DataLoader.
119    """
120    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
121    dataset = get_jnuifm_dataset(path, patch_shape, resize_inputs, download, **ds_kwargs)
122    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)