torch_em.data.datasets.medical.thyrln_pumch

  1"""The ThyRLN-PUMCH dataset contains annotations for recurrent laryngeal nerve (RLN) segmentation
  2in endoscopic thyroidectomy video frames.
  3
  4The dataset comprises frames extracted from 28 endoscopic thyroidectomy surgeries performed at
  5Peking Union Medical College Hospital (PUMCH), with 18,178 pixel-level RLN segmentation masks
  6(binary masks, manually annotated by two endocrine surgeons and reviewed / modified by senior
  7head-and-neck surgeons), plus 734 explicitly negative frames (no RLN visible). Annotations are
  8provided as per-frame binary masks paired with the raw jpg frames.
  9
 10The dataset is located at https://www.kaggle.com/datasets/zhenghuaijin/thyrln-pumch and is
 11distributed under the CC0 1.0 license. Downloading it requires the Kaggle API and a Kaggle
 12account with an API token, see `torch_em.data.datasets.util.download_source_kaggle` for details.
 13
 14This dataset is from the publication https://doi.org/10.1038/s41597-026-06961-6.
 15Please cite it if you use this dataset in your research.
 16"""
 17
 18import os
 19from natsort import natsorted
 20from typing import Union, Tuple, List
 21
 22from torch.utils.data import Dataset, DataLoader
 23
 24import torch_em
 25
 26from .. import util
 27
 28
 29KAGGLE_DATASET_NAME = "zhenghuaijin/thyrln-pumch"
 30
 31
 32def get_thyrln_pumch_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 33    """Download the ThyRLN-PUMCH data.
 34
 35    Args:
 36        path: Filepath to a folder where the data is downloaded for further processing.
 37        download: Whether to download the data if it is not present.
 38
 39    Returns:
 40        Filepath to the extracted ThyRLN-PUMCH data.
 41    """
 42    data_dir = os.path.join(path, "ThyRLN-PUMCH")
 43    if os.path.exists(data_dir):
 44        return data_dir
 45
 46    os.makedirs(path, exist_ok=True)
 47
 48    util.download_source_kaggle(path=path, dataset_name=KAGGLE_DATASET_NAME, download=download)
 49
 50    zip_path = os.path.join(path, "thyrln-pumch.zip")
 51    util.unzip(zip_path=zip_path, dst=path)
 52
 53    if not os.path.exists(data_dir):
 54        raise RuntimeError(f"The dataset could not be found at '{data_dir}' after extraction.")
 55
 56    return data_dir
 57
 58
 59def _is_junk(fname):
 60    return fname.startswith("._")
 61
 62
 63def _strip_mp4_suffix(name):
 64    return name[:-4] if name.lower().endswith(".mp4") else name
 65
 66
 67def get_thyrln_pumch_paths(path: Union[os.PathLike, str], download: bool = False) -> Tuple[List[str], List[str]]:
 68    """Get paths to the ThyRLN-PUMCH data.
 69
 70    This returns the frames with an explicit pixel-level RLN mask ('labeled picture'). Frames are
 71    matched to their raw jpg either from the 'labeled picture' or, as a fallback, the (more complete)
 72    'original picture' folder, since not every annotated frame is duplicated into 'labeled picture'.
 73
 74    Args:
 75        path: Filepath to a folder where the data is downloaded for further processing.
 76        download: Whether to download the data if it is not present.
 77
 78    Returns:
 79        List of filepaths for the image data.
 80        List of filepaths for the label data.
 81    """
 82    data_dir = get_thyrln_pumch_data(path, download)
 83
 84    labeled_dir = os.path.join(data_dir, "labeled picture")
 85    original_dir = os.path.join(data_dir, "original picture")
 86
 87    image_paths, mask_paths = [], []
 88    for surgery in natsorted(os.listdir(labeled_dir)):
 89        surgery_dir = os.path.join(labeled_dir, surgery)
 90        if not os.path.isdir(surgery_dir):
 91            continue
 92
 93        mask_dir = os.path.join(surgery_dir, "mask")
 94        if not os.path.isdir(mask_dir):
 95            continue
 96
 97        labeled_video_dirs = {
 98            _strip_mp4_suffix(d): d for d in os.listdir(surgery_dir)
 99            if d != "mask" and os.path.isdir(os.path.join(surgery_dir, d))
100        }
101
102        original_surgery_dir = os.path.join(original_dir, surgery)
103        original_video_dirs = {
104            _strip_mp4_suffix(d): d for d in os.listdir(original_surgery_dir)
105        } if os.path.isdir(original_surgery_dir) else {}
106
107        for mask_video_dir in natsorted(os.listdir(mask_dir)):
108            key = _strip_mp4_suffix(mask_video_dir)
109            mask_video_path = os.path.join(mask_dir, mask_video_dir)
110            if not os.path.isdir(mask_video_path):
111                continue
112
113            labeled_img_dir = os.path.join(surgery_dir, labeled_video_dirs.get(key, ""))
114            original_img_dir = os.path.join(original_surgery_dir, original_video_dirs.get(key, ""))
115
116            for mfname in natsorted(os.listdir(mask_video_path)):
117                if _is_junk(mfname) or not mfname.endswith("_mask.png"):
118                    continue
119
120                stem = mfname[:-len("_mask.png")]
121                image_fname = f"{stem}.jpg"
122
123                image_path = os.path.join(labeled_img_dir, image_fname)
124                if not os.path.exists(image_path):
125                    image_path = os.path.join(original_img_dir, image_fname)
126                if not os.path.exists(image_path):
127                    continue
128
129                image_paths.append(image_path)
130                mask_paths.append(os.path.join(mask_video_path, mfname))
131
132    if len(image_paths) == 0 or len(image_paths) != len(mask_paths):
133        raise RuntimeError("Something went wrong with fetching the image and label paths.")
134
135    return image_paths, mask_paths
136
137
138def get_thyrln_pumch_dataset(
139    path: Union[os.PathLike, str],
140    patch_shape: Tuple[int, int],
141    resize_inputs: bool = False,
142    download: bool = False,
143    **kwargs
144) -> Dataset:
145    """Get the ThyRLN-PUMCH dataset for recurrent laryngeal nerve segmentation.
146
147    Args:
148        path: Filepath to a folder where the data is downloaded for further processing.
149        patch_shape: The patch shape to use for training.
150        resize_inputs: Whether to resize the inputs.
151        download: Whether to download the data if it is not present.
152        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
153
154    Returns:
155        The segmentation dataset.
156    """
157    image_paths, mask_paths = get_thyrln_pumch_paths(path, download)
158
159    if resize_inputs:
160        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": True}
161        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
162            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
163        )
164
165    return torch_em.default_segmentation_dataset(
166        raw_paths=image_paths,
167        raw_key=None,
168        label_paths=mask_paths,
169        label_key=None,
170        patch_shape=patch_shape,
171        is_seg_dataset=False,
172        **kwargs
173    )
174
175
176def get_thyrln_pumch_loader(
177    path: Union[os.PathLike, str],
178    batch_size: int,
179    patch_shape: Tuple[int, int],
180    resize_inputs: bool = False,
181    download: bool = False,
182    **kwargs
183) -> DataLoader:
184    """Get the ThyRLN-PUMCH dataloader for recurrent laryngeal nerve segmentation.
185
186    Args:
187        path: Filepath to a folder where the data is downloaded for further processing.
188        batch_size: The batch size for training.
189        patch_shape: The patch shape to use for training.
190        resize_inputs: Whether to resize the inputs.
191        download: Whether to download the data if it is not present.
192        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
193
194    Returns:
195        The DataLoader.
196    """
197    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
198    dataset = get_thyrln_pumch_dataset(path, patch_shape, resize_inputs, download, **ds_kwargs)
199    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)