torch_em.data.datasets.medical.kvasir_instrument

  1"""The Kvasir-Instrument dataset contains annotations for surgical instrument
  2(e.g. snares, biopsy forceps) segmentation in gastrointestinal endoscopy.
  3
  4NOTE: This is a different dataset from Kvasir-SEG (polyp segmentation, see `kvasir.py`).
  5
  6The dataset is located at https://datasets.simula.no/kvasir-instrument/,
  7mirrored on Kaggle at https://www.kaggle.com/datasets/debeshjha1/kvasirinstrument.
  8This dataset is from the publication https://doi.org/10.1007/978-3-030-67835-7_19.
  9Please cite it if you use this dataset for your research.
 10"""
 11
 12import os
 13from tqdm import tqdm
 14from natsort import natsorted
 15from typing import Union, Tuple, List, Literal
 16
 17import numpy as np
 18import imageio.v3 as imageio
 19
 20from torch.utils.data import Dataset, DataLoader
 21
 22import torch_em
 23
 24from .. import util
 25
 26
 27CHECKSUM = "8e12d6c9e232e2f3db90f325ae6e5f81143b5704afcf94d94835c6aa3f1314e1"
 28
 29
 30def get_kvasir_instrument_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 31    """Download the Kvasir-Instrument dataset.
 32
 33    Args:
 34        path: Filepath to a folder where the data is downloaded for further processing.
 35        download: Whether to download the data if it is not present.
 36
 37    Returns:
 38        Filepath where the data is downloaded.
 39    """
 40    data_dir = os.path.join(path, "kvasir-instrument")
 41    if os.path.exists(data_dir):
 42        return data_dir
 43
 44    os.makedirs(path, exist_ok=True)
 45
 46    util.download_source_kaggle(path=path, dataset_name="debeshjha1/kvasirinstrument", download=download)
 47    zip_path = os.path.join(path, "kvasirinstrument.zip")
 48    util._check_checksum(zip_path, CHECKSUM)
 49    util.unzip(zip_path=zip_path, dst=path)
 50
 51    return data_dir
 52
 53
 54def get_kvasir_instrument_paths(
 55    path: Union[os.PathLike, str], split: Literal["train", "test"], download: bool = False
 56) -> Tuple[List[str], List[str]]:
 57    """Get paths to the Kvasir-Instrument data.
 58
 59    Args:
 60        path: Filepath to a folder where the data is downloaded for further processing.
 61        split: The choice of data split.
 62        download: Whether to download the data if it is not present.
 63
 64    Returns:
 65        List of filepaths for the image data.
 66        List of filepaths for the label data.
 67    """
 68    data_dir = get_kvasir_instrument_data(path=path, download=download)
 69
 70    split_file = os.path.join(data_dir, f"{split}.txt")
 71    with open(split_file) as f:
 72        image_ids = [line.strip() for line in f if line.strip()]
 73
 74    neu_gt_dir = os.path.join(data_dir, "masks", "preprocessed")
 75    os.makedirs(neu_gt_dir, exist_ok=True)
 76
 77    image_paths, gt_paths = [], []
 78    for image_id in tqdm(image_ids):
 79        image_path = os.path.join(data_dir, "images", "images", f"{image_id}.jpg")
 80        gt_path = os.path.join(data_dir, "masks", "masks", f"{image_id}.png")
 81        neu_gt_path = os.path.join(neu_gt_dir, f"{image_id}.tif")
 82
 83        image_paths.append(image_path)
 84        gt_paths.append(neu_gt_path)
 85        if os.path.exists(neu_gt_path):
 86            continue
 87
 88        gt = imageio.imread(gt_path)
 89        gt = np.mean(gt, axis=-1)
 90        gt = (gt >= 240).astype("uint8")
 91        imageio.imwrite(neu_gt_path, gt, compression="zlib")
 92
 93    return natsorted(image_paths), natsorted(gt_paths)
 94
 95
 96def get_kvasir_instrument_dataset(
 97    path: Union[os.PathLike, str],
 98    patch_shape: Tuple[int, int],
 99    split: Literal["train", "test"],
100    resize_inputs: bool = False,
101    download: bool = False,
102    **kwargs
103) -> Dataset:
104    """Get the Kvasir-Instrument dataset for surgical instrument segmentation.
105
106    Args:
107        path: Filepath to a folder where the data is downloaded for further processing.
108        patch_shape: The patch shape to use for training.
109        split: The choice of data split.
110        resize_inputs: Whether to resize the inputs to the patch shape.
111        download: Whether to download the data if it is not present.
112        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
113
114    Returns:
115        The segmentation dataset.
116    """
117    image_paths, gt_paths = get_kvasir_instrument_paths(path, split, download)
118
119    if resize_inputs:
120        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": True}
121        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
122            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
123        )
124
125    return torch_em.default_segmentation_dataset(
126        raw_paths=image_paths,
127        raw_key=None,
128        label_paths=gt_paths,
129        label_key=None,
130        patch_shape=patch_shape,
131        is_seg_dataset=False,
132        **kwargs
133    )
134
135
136def get_kvasir_instrument_loader(
137    path: Union[os.PathLike, str],
138    patch_shape: Tuple[int, int],
139    batch_size: int,
140    split: Literal["train", "test"],
141    resize_inputs: bool = False,
142    download: bool = False,
143    **kwargs
144) -> DataLoader:
145    """Get the Kvasir-Instrument dataloader for surgical instrument 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        batch_size: The batch size for training.
151        split: The choice of data split.
152        resize_inputs: Whether to resize the inputs to the patch shape.
153        download: Whether to download the data if it is not present.
154        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
155
156    Returns:
157        The DataLoader.
158    """
159    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
160    dataset = get_kvasir_instrument_dataset(path, patch_shape, split, resize_inputs, download, **ds_kwargs)
161    return torch_em.get_data_loader(dataset=dataset, batch_size=batch_size, **loader_kwargs)