torch_em.data.datasets.medical.spine_endoscopic_atlas

  1"""The Spine Endoscopic Atlas (SEA) dataset contains annotations for surgical instrument
  2segmentation in endoscopic spine surgery images.
  3
  4The full archive ships 48,510 images in total, of which 10,662 come with an instrument
  5segmentation mask (stored as an NRRD file in a sibling 'seg' folder, matched to its image by
  6basename); the remaining ~37,848 images are unlabeled raw frames under 'unclassified' and are
  7not used by this loader. Annotated images are organized by working-channel diameter ('big' or
  8'small'), spinal region ('cervical' or 'lumbar'), patient, difficulty ('normal' or 'difficult'
  9scenario) and instrument type ('bipolar', 'grasping_forceps', 'drill', 'dissector', 'punch' or
 10'scissor'); the instrument type can be selected with the 'instrument' argument. This loader
 11discovers the image-mask pairs on disk rather than relying on the advertised counts, and merges
 12each NRRD segmentation (which may encode either a binary 0/255 mask or a per-segment integer
 13label map, depending on the file) into a binary instrument mask during preprocessing.
 14
 15The dataset is located at https://doi.org/10.6084/m9.figshare.27109312, released under a CC0 license.
 16"""
 17
 18import os
 19from glob import glob
 20from natsort import natsorted
 21from typing import Union, Tuple, List, Optional
 22
 23from torch.utils.data import Dataset, DataLoader
 24
 25import torch_em
 26
 27from .. import util
 28
 29
 30URL = "https://ndownloader.figshare.com/files/56202515"
 31CHECKSUM = "bde2014aa786181140552ab248fad0b2"
 32
 33INSTRUMENTS = ["bipolar", "grasping_forceps", "drill", "dissector", "punch", "scissor"]
 34
 35
 36def _convert_mask(nrrd_path, out_path):
 37    import numpy as np
 38    import SimpleITK as sitk
 39    import tifffile
 40
 41    if os.path.exists(out_path):
 42        return
 43
 44    data = sitk.GetArrayFromImage(sitk.ReadImage(nrrd_path))
 45    data = np.squeeze(data)
 46    mask = (data > 0).astype("uint8")
 47
 48    tmp_path = f"{out_path}.{os.getpid()}.incomplete.tif"
 49    tifffile.imwrite(tmp_path, mask)
 50    os.replace(tmp_path, out_path)
 51
 52
 53def _convert_raw(raw_path, out_path):
 54    from PIL import Image
 55
 56    if os.path.exists(out_path):
 57        return
 58
 59    image = Image.open(raw_path).convert("RGB")
 60
 61    tmp_path = f"{out_path}.{os.getpid()}.incomplete.png"
 62    image.save(tmp_path)
 63    os.replace(tmp_path, out_path)
 64
 65
 66def get_spine_endoscopic_atlas_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 67    """Download the Spine Endoscopic Atlas dataset.
 68
 69    Args:
 70        path: Filepath to a folder where the data is downloaded for further processing.
 71        download: Whether to download the data if it is not present.
 72
 73    Returns:
 74        Filepath where the data is downloaded.
 75    """
 76    data_dir = os.path.join(path, "Spine endoscopic atlas", "classified")
 77    if os.path.exists(data_dir):
 78        return data_dir
 79
 80    os.makedirs(path, exist_ok=True)
 81
 82    zip_path = os.path.join(path, "spine_endoscopic_atlas.zip")
 83    util.download_source(path=zip_path, url=URL, download=download, checksum=CHECKSUM)
 84    util.unzip(zip_path=zip_path, dst=path, remove=False)
 85
 86    assert os.path.exists(data_dir), f"The extraction of the archive did not create the expected folder in '{path}'."
 87
 88    return data_dir
 89
 90
 91def get_spine_endoscopic_atlas_paths(
 92    path: Union[os.PathLike, str], instrument: Optional[str] = None, download: bool = False,
 93) -> Tuple[List[str], List[str]]:
 94    """Get paths to the Spine Endoscopic Atlas data.
 95
 96    Args:
 97        path: Filepath to a folder where the data is downloaded for further processing.
 98        instrument: The choice of instrument type to restrict the data to. By default all instrument
 99            types are used. See `INSTRUMENTS` for the valid choices.
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    if instrument is not None and instrument not in INSTRUMENTS:
107        raise ValueError(f"'{instrument}' is not a valid instrument. Choose one of {INSTRUMENTS}.")
108
109    data_dir = get_spine_endoscopic_atlas_data(path, download)
110    mask_dir = os.path.join(path, "spine_endoscopic_atlas_masks")
111    raw_dir = os.path.join(path, "spine_endoscopic_atlas_raw")
112    os.makedirs(mask_dir, exist_ok=True)
113    os.makedirs(raw_dir, exist_ok=True)
114
115    pattern = os.path.join(data_dir, "*", "*", "*", "*", instrument or "*", "*.*")
116    raw_paths, label_paths = [], []
117    for source_path in natsorted(glob(pattern)):
118        if source_path.lower().endswith((".jpg", ".png")) and os.path.basename(os.path.dirname(source_path)) != "seg":
119            base = os.path.splitext(os.path.basename(source_path))[0]
120            nrrd_path = os.path.join(os.path.dirname(source_path), "seg", f"{base}.nrrd")
121            if not os.path.exists(nrrd_path):
122                continue
123
124            # The raw images are re-encoded as RGB PNGs: some ship as RGBA, which breaks the
125            # fixed 3-channel assumption of the resize transform used by 'resize_inputs'.
126            raw_path = os.path.join(raw_dir, f"{base}.png")
127            _convert_raw(source_path, raw_path)
128
129            label_path = os.path.join(mask_dir, f"{base}.tif")
130            _convert_mask(nrrd_path, label_path)
131
132            raw_paths.append(raw_path)
133            label_paths.append(label_path)
134
135    assert len(raw_paths) == len(label_paths) and len(raw_paths) > 0
136
137    return raw_paths, label_paths
138
139
140def get_spine_endoscopic_atlas_dataset(
141    path: Union[os.PathLike, str],
142    patch_shape: Tuple[int, int],
143    instrument: Optional[str] = None,
144    resize_inputs: bool = False,
145    download: bool = False,
146    **kwargs
147) -> Dataset:
148    """Get the Spine Endoscopic Atlas dataset for surgical instrument segmentation.
149
150    Args:
151        path: Filepath to a folder where the data is downloaded for further processing.
152        patch_shape: The patch shape to use for training.
153        instrument: The choice of instrument type to restrict the data to.
154        resize_inputs: Whether to resize the inputs to the patch shape.
155        download: Whether to download the data if it is not present.
156        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
157
158    Returns:
159        The segmentation dataset.
160    """
161    raw_paths, label_paths = get_spine_endoscopic_atlas_paths(path, instrument, download)
162
163    if resize_inputs:
164        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": True}
165        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
166            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
167        )
168
169    return torch_em.default_segmentation_dataset(
170        raw_paths=raw_paths,
171        raw_key=None,
172        label_paths=label_paths,
173        label_key=None,
174        patch_shape=patch_shape,
175        is_seg_dataset=False,
176        **kwargs
177    )
178
179
180def get_spine_endoscopic_atlas_loader(
181    path: Union[os.PathLike, str],
182    batch_size: int,
183    patch_shape: Tuple[int, int],
184    instrument: Optional[str] = None,
185    resize_inputs: bool = False,
186    download: bool = False,
187    **kwargs
188) -> DataLoader:
189    """Get the Spine Endoscopic Atlas dataloader for surgical instrument segmentation.
190
191    Args:
192        path: Filepath to a folder where the data is downloaded for further processing.
193        batch_size: The batch size for training.
194        patch_shape: The patch shape to use for training.
195        instrument: The choice of instrument type to restrict the data to.
196        resize_inputs: Whether to resize the inputs to the patch shape.
197        download: Whether to download the data if it is not present.
198        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
199
200    Returns:
201        The DataLoader.
202    """
203    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
204    dataset = get_spine_endoscopic_atlas_dataset(path, patch_shape, instrument, resize_inputs, download, **ds_kwargs)
205    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)