torch_em.data.datasets.medical.m2caiseg

  1"""The m2caiseg dataset contains annotations for segmentation of organs and instruments in endoscopy.
  2
  3This dataset is located at https://www.kaggle.com/datasets/salmanmaq/m2caiseg.
  4The data is from the publication https://doi.org/10.48550/arXiv.2008.10134.
  5Please cite it if you use this data in a publication.
  6"""
  7
  8import os
  9from glob import glob
 10from tqdm import tqdm
 11from pathlib import Path
 12from natsort import natsorted
 13from typing import Union, Tuple, List, Literal
 14
 15import numpy as np
 16import imageio.v3 as imageio
 17
 18from torch.utils.data import Dataset, DataLoader
 19
 20import torch_em
 21
 22from .. import util
 23
 24
 25LABEL_MAPS = {
 26    (0, 0, 0): 0,  # out of frame
 27    (0, 85, 170): 1,  # grasper
 28    (0, 85, 255): 2,  # bipolar
 29    (0, 170, 255): 3,  # hook
 30    (0, 255, 85): 4,  # scissors
 31    (0, 255, 170): 5,  # clipper
 32    (85, 0, 170): 6,  # irrigator
 33    (85, 0, 255): 7,  # specimen bag
 34    (170, 85, 85): 8,  # trocars
 35    (170, 170, 170): 9,  # clip
 36    (85, 170, 0): 10,  # liver
 37    (85, 170, 255): 11,  # gall bladder
 38    (85, 255, 0): 12,  # fat
 39    (85, 255, 170): 13,  # upper wall
 40    (170, 0, 255): 14,  # artery
 41    (255, 0, 255): 15,  # intestine
 42    (255, 255, 0): 16,  # bile
 43    (255, 0, 0): 17,  # blood
 44    (170, 0, 85): 18,  # unknown
 45}
 46
 47
 48def get_m2caiseg_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 49    """Get the m2caiseg dataset.
 50
 51    Args:
 52        path: Filepath to a folder where the data is downloaded for further processing.
 53        download: Whether to download the data if it is not present.
 54
 55    Returns:
 56        Filepath where the data is downloaded.
 57    """
 58    data_dir = os.path.join(path, r"m2caiSeg dataset")
 59    if os.path.exists(data_dir):
 60        return data_dir
 61
 62    os.makedirs(path, exist_ok=True)
 63
 64    util.download_source_kaggle(path=path, dataset_name="salmanmaq/m2caiseg", download=download)
 65    zip_path = os.path.join(path, "m2caiseg.zip")
 66    util.unzip(zip_path=zip_path, dst=path)
 67
 68    return data_dir
 69
 70
 71def get_m2caiseg_paths(
 72    path: Union[os.PathLike, str], split: Literal["train", "val", "test"], download: bool = False
 73) -> Tuple[List[str], List[str]]:
 74    """Get paths to the m2caiseg 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.
 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    data_dir = get_m2caiseg_data(path=path, download=download)
 86
 87    if split == "val":
 88        impaths = natsorted(glob(os.path.join(data_dir, "train", "images", "*.jpg")))
 89        gpaths = natsorted(glob(os.path.join(data_dir, "train", "groundtruth", "*.png")))
 90
 91        imids = [os.path.split(_p)[-1] for _p in impaths]
 92        gids = [os.path.split(_p)[-1] for _p in gpaths]
 93
 94        image_paths = [
 95            _p for _p in natsorted(
 96                glob(os.path.join(data_dir, "trainval", "images", "*.jpg"))
 97            ) if os.path.split(_p)[-1] not in imids
 98        ]
 99        gt_paths = [
100            _p for _p in natsorted(
101                glob(os.path.join(data_dir, "trainval", "groundtruth", "*.png"))
102            ) if os.path.split(_p)[-1] not in gids
103        ]
104
105    else:
106        image_paths = natsorted(glob(os.path.join(data_dir, split, "images", "*.jpg")))
107        gt_paths = natsorted(glob(os.path.join(data_dir, split, "groundtruth", "*.png")))
108
109    images_dir = os.path.join(data_dir, "preprocessed", split, "images")
110    mask_dir = os.path.join(data_dir, "preprocessed", split, "masks")
111    if os.path.exists(images_dir) and os.path.exists(mask_dir):
112        return natsorted(glob(os.path.join(images_dir, "*"))), natsorted(glob(os.path.join(mask_dir, "*")))
113
114    os.makedirs(images_dir, exist_ok=True)
115    os.makedirs(mask_dir, exist_ok=True)
116
117    fimage_paths, fgt_paths = [], []
118    for image_path, gt_path in tqdm(zip(image_paths, gt_paths), total=len(image_paths)):
119        image = imageio.imread(image_path)
120        gt = imageio.imread(gt_path)
121
122        image_id = Path(image_path).stem
123        gt_id = Path(gt_path).stem
124
125        if image.shape != gt.shape:
126            print("This pair of image and labels mismatch.")
127            continue
128
129        dst_image_path = os.path.join(images_dir, f"{image_id}.tif")
130        dst_gt_path = os.path.join(mask_dir, f"{gt_id}.tif")
131
132        fimage_paths.append(image_path)
133        fgt_paths.append(dst_gt_path)
134        if os.path.exists(dst_gt_path) and os.path.exists(dst_image_path):
135            continue
136
137        instances = np.zeros(gt.shape[:2])
138        for lmap in LABEL_MAPS:
139            binary_map = (gt == lmap).all(axis=2)
140            instances[binary_map > 0] = LABEL_MAPS[lmap]
141
142        imageio.imwrite(dst_image_path, image, compression="zlib")
143        imageio.imwrite(dst_gt_path, instances, compression="zlib")
144
145    return fimage_paths, fgt_paths
146
147
148def get_m2caiseg_dataset(
149    path: Union[os.PathLike, str],
150    patch_shape: Tuple[int, int],
151    split: Literal["train", "val", "test"],
152    resize_inputs: bool = False,
153    download: bool = False,
154    **kwargs
155) -> Dataset:
156    """Get the m2caiseg dataset for organ and instrument segmentation.
157
158    Args:
159        path: Filepath to a folder where the data is downloaded for further processing.
160        patch_shape: The patch shape to use for training.
161        split: The choice of data split.
162        resize_inputs: Whether to resize inputs to the desired patch shape.
163        download: Whether to download the data if it is not present.
164        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
165
166    Returns:
167        The segmentation dataset.
168    """
169    image_paths, gt_paths = get_m2caiseg_paths(path, split, download)
170
171    if resize_inputs:
172        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": True}
173        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
174            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
175        )
176
177    return torch_em.default_segmentation_dataset(
178        raw_paths=image_paths,
179        raw_key=None,
180        label_paths=gt_paths,
181        label_key=None,
182        patch_shape=patch_shape,
183        is_seg_dataset=False,
184        **kwargs
185    )
186
187
188def get_m2caiseg_loader(
189    path: Union[os.PathLike, str],
190    batch_size: int,
191    patch_shape: Tuple[int, int],
192    split: Literal["train", "val", "test"],
193    resize_inputs: bool = False,
194    download: bool = False,
195    **kwargs
196) -> DataLoader:
197    """Get the m2caiseg dataloader for organ and instrument segmentation.
198
199    Args:
200        path: Filepath to a folder where the data is downloaded for further processing.
201        batch_size: The batch size for training.
202        patch_shape: The patch shape to use for training.
203        split: The choice of data split.
204        resize_inputs: Whether to resize inputs to the desired patch shape.
205        download: Whether to download the data if it is not present.
206        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
207
208    Returns:
209        The DataLoader.
210    """
211    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
212    dataset = get_m2caiseg_dataset(path, patch_shape, split, resize_inputs, download, **ds_kwargs)
213    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)