torch_em.data.datasets.medical.lumbar_spine_us

  1"""The Lumbar Spine US dataset contains annotations for lumbar bone surface segmentation
  2in paired handheld (HUS) and robot-assisted (RUS) ultrasound frames, acquired together with
  3ground-truth CT of the lumbar spine in 63 healthy volunteers. Out of these, 9 participants
  4have expert-annotated bone surface masks for 6091 ultrasound frames in total (2353 HUS and
  53738 RUS frames, according to the publication below).
  6
  7The dataset is located at https://doi.org/10.48804/3XPCAE.
  8This dataset is from the publication https://doi.org/10.1038/s41597-025-06047-9.
  9Please cite it if you use this dataset in your research.
 10"""
 11
 12import os
 13from glob import glob
 14from tqdm import tqdm
 15from natsort import natsorted
 16from typing import Union, Tuple, Literal, List
 17
 18import imageio.v3 as imageio
 19
 20from torch.utils.data import Dataset, DataLoader
 21
 22import torch_em
 23
 24from .. import util
 25
 26
 27BASE_URL = "https://rdr.kuleuven.be/api/access/datafile/"
 28
 29# Each entry maps a "<participant>_<scan>" recording, hosted on the KU Leuven RDR (Dataverse)
 30# repository, to its (file id, sha256 checksum, probe). Scans tagged 'H' are handheld (HUS)
 31# acquisitions. Scans tagged 'R' and 'D' are robot-assisted (RUS) acquisitions: 'D' denotes the
 32# robotic scan types other than the 'Perpendicular' and 'along the spinous process' ones (see the
 33# publication's Data Records section for the scan type naming convention).
 34SCANS = {
 35    "URS08_H1": (226877, "5b68e478100e5c445bbcabb8df67e6ea2aa7438e7ef85ac9445f3c93f1ebab8e", "handheld"),
 36    "URS08_R2": (226871, "62a1c0abcab00b894009ea2f7a06cfecfbaf95421b896cfbf65492797fcc632b", "robotic"),
 37    "URS16_H3": (227333, "3cf4892554902c724a0008453369eda2986a41b38c87571e53f0373009b5753d", "handheld"),
 38    "URS16_R1": (227210, "23d0738403c5ccf6bf895b9b363e758011804cf1709fcdf59791dd041c8eabd9", "robotic"),
 39    "URS26_H3": (226883, "0fbf944aa9acadc2bcb244d4e1234fd7330ad776ef7ba8cae856d3b62e251f92", "handheld"),
 40    "URS26_R1": (226873, "ba52bef856bc34e4c91177e976ca79982f22feb058a40eac009cf8c87159dca8", "robotic"),
 41    "URS31_D2": (227069, "eb18d445d1612cca6f243d9a6dd71f0b53bc40827569b69daf1e08d17e51227a", "robotic"),
 42    "URS31_R2": (226964, "ac8a69f34b9c2eff28ded73e2ddb5b44b15ae9bec240ea1eef23a5335c587733", "robotic"),
 43    "URS36_D2": (226975, "ffdcfef013706d43bf03011efc121507dca4ab3adfa771e697697a65d27bb29a", "robotic"),
 44    "URS36_H2": (227136, "5e203254f081ec5559949aef601675fef49f50b1d8c307d738c2d57e76621a69", "handheld"),
 45    "URS40_H4": (227343, "67fadae809baca2478858fecd8074acbdb0493501f4b3772e78669a0ad6ded6a", "handheld"),
 46    "URS40_R1": (227007, "430ea76b2eefe038f3e77c006417e597a4c6522830c367461cb07c46b55611a6", "robotic"),
 47    "URS45_H2": (227396, "3fb13ff323ef7cb6e8fe519bdb24784a3eeddbe74a0938b9eb4c46911aace470", "handheld"),
 48    "URS45_R2": (227066, "bcb4e83340b0edeec6e6a061d1d455702e49e60894bfec60a5b0b2702fca1046", "robotic"),
 49    "URS51_D2": (226893, "8d2e687c769f4fac6ce22a3acc51c2bedf79636b56fbe0bc9ec8ff0d7bf2c0df", "robotic"),
 50    "URS51_R2": (226825, "ac89115922ee6bf54a5c5fbc1710755e8d2386b2fb69caad67efc1e9c1b28ded", "robotic"),
 51    "URS54_D1": (226951, "8e4b8f45a0a9dc149a9f78ba6eb88e7ef817fe0535815607c6db862ab8c7ebcb", "robotic"),
 52    "URS54_H3": (226958, "9037e202700de1ce3f181a3937b9bbd1e43499526545ae7388375f489d04ba97", "handheld"),
 53}
 54
 55
 56def _read_metaimage(path):
 57    """Read a MetaImage ('.mhd' + '.raw') frame, returning it as a numpy array. Supports zlib compressed
 58    data, as used by the raw files shipped with this dataset.
 59    """
 60    import SimpleITK as sitk
 61    return sitk.GetArrayFromImage(sitk.ReadImage(path))
 62
 63
 64def get_lumbar_spine_us_data(
 65    path: Union[os.PathLike, str], probe: Literal["handheld", "robotic", "all"] = "all", download: bool = False
 66) -> str:
 67    """Download the Lumbar Spine US dataset.
 68
 69    Args:
 70        path: Filepath to a folder where the data is downloaded for further processing.
 71        probe: The choice of ultrasound probe. Either 'handheld', 'robotic' or 'all'.
 72        download: Whether to download the data if it is not present.
 73
 74    Returns:
 75        Filepath where the data is downloaded.
 76    """
 77    os.makedirs(path, exist_ok=True)
 78
 79    scans = [name for name, (_, _, p) in SCANS.items() if probe == "all" or p == probe]
 80    for name in scans:
 81        scan_dir = os.path.join(path, name)
 82        if os.path.exists(scan_dir):
 83            continue
 84
 85        file_id, checksum, _ = SCANS[name]
 86        zip_path = os.path.join(path, f"{name}.zip")
 87        util.download_source(path=zip_path, url=f"{BASE_URL}{file_id}", download=download, checksum=checksum)
 88        util.unzip(zip_path=zip_path, dst=path)
 89
 90    return path
 91
 92
 93def _convert_scans_to_tif(path, scans):
 94    converted_dir = os.path.join(path, "converted")
 95    os.makedirs(converted_dir, exist_ok=True)
 96
 97    image_paths, label_paths = [], []
 98    for name in scans:
 99        label_mhds = natsorted(glob(os.path.join(path, name, "Labels", "*-labels.mhd")))
100        for label_mhd in tqdm(label_mhds, desc=f"Converting '{name}' to tif", leave=False):
101            frame_id = os.path.basename(label_mhd).replace("-labels.mhd", "")
102            raw_mhd = os.path.join(path, name, "Labels", f"{frame_id}.mhd")
103
104            image_path = os.path.join(converted_dir, f"{name}_{frame_id}.tif")
105            label_path = os.path.join(converted_dir, f"{name}_{frame_id}_labels.tif")
106            image_paths.append(image_path)
107            label_paths.append(label_path)
108            if os.path.exists(image_path) and os.path.exists(label_path):
109                continue
110
111            image = _read_metaimage(raw_mhd)
112            labels = _read_metaimage(label_mhd).astype("uint8")
113            imageio.imwrite(image_path, image)
114            imageio.imwrite(label_path, labels)
115
116    return image_paths, label_paths
117
118
119def get_lumbar_spine_us_paths(
120    path: Union[os.PathLike, str], probe: Literal["handheld", "robotic", "all"] = "all", download: bool = False
121) -> Tuple[List[str], List[str]]:
122    """Get paths to the Lumbar Spine US data.
123
124    Args:
125        path: Filepath to a folder where the data is downloaded for further processing.
126        probe: The choice of ultrasound probe. Either 'handheld', 'robotic' or 'all'.
127        download: Whether to download the data if it is not present.
128
129    Returns:
130        List of filepaths for the image data.
131        List of filepaths for the label data.
132    """
133    get_lumbar_spine_us_data(path=path, probe=probe, download=download)
134
135    scans = [name for name, (_, _, p) in SCANS.items() if probe == "all" or p == probe]
136    image_paths, label_paths = _convert_scans_to_tif(path, scans)
137
138    assert len(image_paths) == len(label_paths) and len(image_paths) > 0
139
140    return image_paths, label_paths
141
142
143def get_lumbar_spine_us_dataset(
144    path: Union[os.PathLike, str],
145    patch_shape: Tuple[int, int],
146    probe: Literal["handheld", "robotic", "all"] = "all",
147    resize_inputs: bool = False,
148    download: bool = False,
149    **kwargs
150) -> Dataset:
151    """Get the Lumbar Spine US dataset for lumbar bone surface segmentation.
152
153    Args:
154        path: Filepath to a folder where the data is downloaded for further processing.
155        patch_shape: The patch shape to use for training.
156        probe: The choice of ultrasound probe. Either 'handheld', 'robotic' or 'all'.
157        resize_inputs: Whether to resize the inputs to the expected patch shape.
158        download: Whether to download the data if it is not present.
159        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
160
161    Returns:
162        The segmentation dataset.
163    """
164    image_paths, label_paths = get_lumbar_spine_us_paths(path=path, probe=probe, download=download)
165
166    if resize_inputs:
167        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
168        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
169            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
170        )
171
172    return torch_em.default_segmentation_dataset(
173        raw_paths=image_paths,
174        raw_key=None,
175        label_paths=label_paths,
176        label_key=None,
177        patch_shape=patch_shape,
178        is_seg_dataset=False,
179        **kwargs
180    )
181
182
183def get_lumbar_spine_us_loader(
184    path: Union[os.PathLike, str],
185    batch_size: int,
186    patch_shape: Tuple[int, int],
187    probe: Literal["handheld", "robotic", "all"] = "all",
188    resize_inputs: bool = False,
189    download: bool = False,
190    **kwargs
191) -> DataLoader:
192    """Get the Lumbar Spine US dataloader for lumbar bone surface segmentation.
193
194    Args:
195        path: Filepath to a folder where the data is downloaded for further processing.
196        batch_size: The batch size for training.
197        patch_shape: The patch shape to use for training.
198        probe: The choice of ultrasound probe. Either 'handheld', 'robotic' or 'all'.
199        resize_inputs: Whether to resize the inputs to the expected patch shape.
200        download: Whether to download the data if it is not present.
201        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
202
203    Returns:
204        The DataLoader.
205    """
206    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
207    dataset = get_lumbar_spine_us_dataset(path, patch_shape, probe, resize_inputs, download, **ds_kwargs)
208    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)