torch_em.data.datasets.medical.robust_mips

  1"""ROBUST-MIPS (Robust Minimally Invasive Pelvic Surgery) is a dataset for surgical
  2instrument instance segmentation and pose estimation in laparoscopic pelvic surgery.
  3
  4The dataset contains 10,040 frames (5,983 for training, 4,057 for testing) sampled from
  5recordings of proctocolectomy, rectal resection and sigmoid resection surgeries. Each frame
  6has a raw endoscopy image, an instance segmentation mask for the surgical instruments, and
  7a JSON file with the instrument tool-tip / pose keypoints (not used by this module).
  8
  9NOTE: The dataset is hosted on Synapse. Downloading it requires the 'synapseclient' python
 10library and a Synapse account with an authentication token stored in the '~/.synapseConfig'
 11file. See 'get_robust_mips_data' for details. The Synapse project 'syn64023381' is public
 12and has no access requirements (confirmed via the Synapse REST API).
 13
 14The dataset is located at https://www.synapse.org/Synapse:syn64023381.
 15This dataset is from the publication https://doi.org/10.48550/arXiv.2508.21096.
 16Please cite it if you use this dataset in your research.
 17"""
 18
 19import os
 20from glob import glob
 21from natsort import natsorted
 22from typing import Union, Tuple, Literal, List
 23
 24from torch.utils.data import Dataset, DataLoader
 25
 26import torch_em
 27
 28from .. import util
 29
 30
 31ENTITY = "syn68915165"
 32
 33
 34def get_robust_mips_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 35    """Download the ROBUST-MIPS dataset.
 36
 37    Follow the instructions below to get access to the dataset.
 38    - Create a free account at https://www.synapse.org.
 39    - Generate a personal access token and store it in a '~/.synapseConfig' file, see
 40      https://python-docs.synapse.org/tutorials/authentication/ for details.
 41    - Install the 'synapseclient' python library.
 42
 43    Args:
 44        path: Filepath to a folder where the data is downloaded for further processing.
 45        download: Whether to download the data if it is not present.
 46
 47    Returns:
 48        Filepath where the data is stored.
 49    """
 50    data_dir = os.path.join(path, "RobustMIPS")
 51    if os.path.exists(data_dir):
 52        return data_dir
 53
 54    os.makedirs(path, exist_ok=True)
 55
 56    import synapseclient
 57
 58    syn = synapseclient.Synapse()
 59    syn.login()
 60    zip_path = os.path.join(path, "RobustMIPS.zip")
 61    if not os.path.exists(zip_path):
 62        if not download:
 63            raise RuntimeError(f"Cannot find the data at {zip_path}, but download was set to False.")
 64        syn.get(ENTITY, downloadLocation=path, downloadFile=True)
 65
 66    util.unzip(zip_path=zip_path, dst=path, remove=False)
 67
 68    return data_dir
 69
 70
 71def get_robust_mips_paths(
 72    path: Union[os.PathLike, str], split: Literal["train", "test"] = "train", download: bool = False
 73) -> Tuple[List[str], List[str]]:
 74    """Get paths to the ROBUST-MIPS data.
 75
 76    Args:
 77        path: Filepath to a folder where the data is downloaded for further processing.
 78        split: The choice of data split. Either 'train' or 'test'.
 79        download: Whether to download the data if it is not present.
 80
 81    Returns:
 82        List of filepaths for the image data.
 83        List of filepaths for the label data.
 84    """
 85    if split not in ["train", "test"]:
 86        raise ValueError(f"'{split}' is not a valid split. Please choose from 'train' or 'test'.")
 87
 88    data_dir = get_robust_mips_data(path, download)
 89
 90    split_dir = "Training" if split == "train" else "Testing"
 91    image_paths = natsorted(glob(os.path.join(data_dir, split_dir, "**", "raw.png"), recursive=True))
 92    gt_paths = [os.path.join(os.path.dirname(p), "instrument_instances.png") for p in image_paths]
 93
 94    assert len(image_paths) > 0, f"No images were found at '{os.path.join(data_dir, split_dir)}'."
 95    assert all(os.path.exists(p) for p in gt_paths), (
 96        "Some 'raw.png' frames do not have a matching 'instrument_instances.png' mask. The expected per-frame "
 97        f"folder layout may not match the actual structure of the downloaded data. Please inspect '{data_dir}'."
 98    )
 99
100    return image_paths, gt_paths
101
102
103def get_robust_mips_dataset(
104    path: Union[os.PathLike, str],
105    patch_shape: Tuple[int, int],
106    split: Literal["train", "test"] = "train",
107    resize_inputs: bool = False,
108    download: bool = False,
109    **kwargs
110) -> Dataset:
111    """Get the ROBUST-MIPS dataset for surgical instrument instance segmentation.
112
113    Args:
114        path: Filepath to a folder where the data is downloaded for further processing.
115        patch_shape: The patch shape to use for training.
116        split: The choice of data split. Either 'train' or 'test'.
117        resize_inputs: Whether to resize inputs to the desired patch shape.
118        download: Whether to download the data if it is not present.
119        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
120
121    Returns:
122        The segmentation dataset.
123    """
124    image_paths, gt_paths = get_robust_mips_paths(path, split, download)
125
126    if resize_inputs:
127        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": True}
128        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
129            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
130        )
131
132    return torch_em.default_segmentation_dataset(
133        raw_paths=image_paths,
134        raw_key=None,
135        label_paths=gt_paths,
136        label_key=None,
137        is_seg_dataset=False,
138        patch_shape=patch_shape,
139        **kwargs
140    )
141
142
143def get_robust_mips_loader(
144    path: Union[os.PathLike, str],
145    batch_size: int,
146    patch_shape: Tuple[int, int],
147    split: Literal["train", "test"] = "train",
148    resize_inputs: bool = False,
149    download: bool = False,
150    **kwargs
151) -> DataLoader:
152    """Get the ROBUST-MIPS dataloader for surgical instrument instance segmentation.
153
154    Args:
155        path: Filepath to a folder where the data is downloaded for further processing.
156        batch_size: The batch size for training.
157        patch_shape: The patch shape to use for training.
158        split: The choice of data split. Either 'train' or 'test'.
159        resize_inputs: Whether to resize inputs to the desired patch shape.
160        download: Whether to download the data if it is not present.
161        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
162
163    Returns:
164        The DataLoader.
165    """
166    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
167    dataset = get_robust_mips_dataset(path, patch_shape, split, resize_inputs, download, **ds_kwargs)
168    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)