torch_em.data.datasets.medical.lss_mri_aisslab

  1"""The LSS MRI AISSLab dataset contains annotations for vertebra, intervertebral disc, sacrum
  2and posterior structure segmentation in sagittal lumbar spine MRI.
  3
  4The dataset consists of 500 patients with sagittal lumbar spine MRI, approved by the IRB of
  5Firat University and clinically validated by neurosurgeons. For each patient, the middle sagittal
  6slice is provided as a PNG image together with a pixel-level segmentation mask that was AI-initialized
  7and then manually refined / verified by the neurosurgeons. The mask labels are (pixel value: label):
  80 = anterior background, 50 = vertebra, 100 = intervertebral disc, 150 = sacrum,
  9200 = posterior A, 255 = posterior B.
 10
 11NOTE: The full release also ships the raw sagittal DICOM series per patient and PNG / XML foraminal
 12stenosis bounding box annotations on all slices. This module only downloads and exposes the
 13middle-slice images and their pixel-level segmentation masks, as the DICOM series and the stenosis
 14bounding boxes are not relevant for segmentation.
 15
 16The dataset is located at https://data.mendeley.com/datasets/rgb77xm3jf/4 (CC BY 4.0).
 17This dataset is from the publication https://doi.org/10.1038/s41597-026-07138-x.
 18Please cite it if you use this dataset for your research.
 19"""
 20
 21import os
 22from glob import glob
 23from natsort import natsorted
 24from typing import Union, Tuple, List
 25
 26import numpy as np
 27import imageio.v3 as imageio
 28
 29from torch.utils.data import Dataset, DataLoader
 30
 31import torch_em
 32
 33from .. import util
 34
 35
 36URL = "https://data.mendeley.com/public-files/datasets/rgb77xm3jf/files/6d9a0116-925d-4111-acb0-1e679f7dfd71/file_downloaded"  # noqa
 37CHECKSUM = "592a294f93d575a16bccc2681c793eb1cfc6679fa2746ac50cbc8970f806b4b1"
 38
 39# The pixel values used in the 'Segmentation/Masks/*M.png' files, mapped to contiguous label ids.
 40LABEL_IDS = {0: 0, 50: 1, 100: 2, 150: 3, 200: 4, 255: 5}
 41
 42
 43def get_lss_mri_aisslab_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 44    """Download the LSS MRI AISSLab dataset.
 45
 46    Args:
 47        path: Filepath to a folder where the data is downloaded for further processing.
 48        download: Whether to download the data if it is not present.
 49
 50    Returns:
 51        Filepath to the folder where the 'Segmentation' data is stored.
 52    """
 53    data_dir = os.path.join(path, "Segmentation")
 54    if os.path.exists(data_dir):
 55        return data_dir
 56
 57    os.makedirs(path, exist_ok=True)
 58
 59    zip_path = os.path.join(path, "LSS_MRI_AISSLab.zip")
 60    util.download_source(path=zip_path, url=URL, download=download, checksum=CHECKSUM)
 61
 62    import zipfile
 63    with zipfile.ZipFile(zip_path) as f:
 64        members = [
 65            m for m in f.namelist()
 66            if m.startswith("Segmentation/Middle_Slice/") or m.startswith("Segmentation/Masks/")
 67        ]
 68        f.extractall(path, members=members)
 69    os.remove(zip_path)
 70
 71    return data_dir
 72
 73
 74def _preprocess_masks(data_dir):
 75    mask_paths = natsorted(glob(os.path.join(data_dir, "Masks", "*M.png")))
 76    neu_dir = os.path.join(data_dir, "preprocessed_masks")
 77    os.makedirs(neu_dir, exist_ok=True)
 78
 79    neu_mask_paths = []
 80    for mask_path in mask_paths:
 81        neu_path = os.path.join(neu_dir, os.path.basename(mask_path))
 82        if not os.path.exists(neu_path):
 83            mask = imageio.imread(mask_path)
 84            neu_mask = np.zeros_like(mask, dtype="uint8")
 85            for value, label_id in LABEL_IDS.items():
 86                neu_mask[mask == value] = label_id
 87            imageio.imwrite(neu_path, neu_mask)
 88        neu_mask_paths.append(neu_path)
 89
 90    return neu_mask_paths
 91
 92
 93def get_lss_mri_aisslab_paths(
 94    path: Union[os.PathLike, str], download: bool = False
 95) -> Tuple[List[str], List[str]]:
 96    """Get paths to the LSS MRI AISSLab data.
 97
 98    Args:
 99        path: Filepath to a folder where the data is downloaded for further processing.
100        download: Whether to download the data if it is not present.
101
102    Returns:
103        List of filepaths for the image data.
104        List of filepaths for the label data.
105    """
106    data_dir = get_lss_mri_aisslab_data(path, download)
107
108    image_paths = natsorted(glob(os.path.join(data_dir, "Middle_Slice", "*.png")))
109    label_paths = _preprocess_masks(data_dir)
110
111    assert len(image_paths) > 0 and len(image_paths) == len(label_paths), (
112        f"Expected the same number of images and masks, got {len(image_paths)} and {len(label_paths)}."
113    )
114
115    return image_paths, label_paths
116
117
118def get_lss_mri_aisslab_dataset(
119    path: Union[os.PathLike, str],
120    patch_shape: Tuple[int, int],
121    resize_inputs: bool = False,
122    download: bool = False,
123    **kwargs
124) -> Dataset:
125    """Get the LSS MRI AISSLab dataset for lumbar spine segmentation in sagittal MRI.
126
127    Args:
128        path: Filepath to a folder where the data is downloaded for further processing.
129        patch_shape: The patch shape to use for training.
130        resize_inputs: Whether to resize the inputs to the expected patch shape.
131        download: Whether to download the data if it is not present.
132        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
133
134    Returns:
135        The segmentation dataset.
136    """
137    image_paths, label_paths = get_lss_mri_aisslab_paths(path, download)
138
139    if resize_inputs:
140        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
141        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
142            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
143        )
144
145    return torch_em.default_segmentation_dataset(
146        raw_paths=image_paths,
147        raw_key=None,
148        label_paths=label_paths,
149        label_key=None,
150        is_seg_dataset=False,
151        patch_shape=patch_shape,
152        **kwargs
153    )
154
155
156def get_lss_mri_aisslab_loader(
157    path: Union[os.PathLike, str],
158    batch_size: int,
159    patch_shape: Tuple[int, int],
160    resize_inputs: bool = False,
161    download: bool = False,
162    **kwargs
163) -> DataLoader:
164    """Get the LSS MRI AISSLab dataloader for lumbar spine segmentation in sagittal MRI.
165
166    Args:
167        path: Filepath to a folder where the data is downloaded for further processing.
168        batch_size: The batch size for training.
169        patch_shape: The patch shape to use for training.
170        resize_inputs: Whether to resize the inputs to the expected patch shape.
171        download: Whether to download the data if it is not present.
172        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
173
174    Returns:
175        The DataLoader.
176    """
177    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
178    dataset = get_lss_mri_aisslab_dataset(path, patch_shape, resize_inputs, download, **ds_kwargs)
179    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)