torch_em.data.datasets.medical.gist514

  1"""The GIST514-DB dataset contains annotations for lesion segmentation in endoscopic ultrasound
  2(EUS) images of gastrointestinal stromal tumors (GISTs) and leiomyomas.
  3
  4The dataset comprises 514 EUS cases (251 GIST, 263 leiomyoma), with lesion contours manually
  5delineated and verified by expert clinicians. It is the dataset introduced for the 'Query2'
  6GIST detection framework.
  7
  8The data is hosted on Google Drive, linked from https://github.com/howardchina/query2, and is
  9distributed under the CC BY-NC-SA 4.0 license.
 10
 11The dataset is from the publication https://doi.org/10.1016/j.compbiomed.2022.106424.
 12Please cite it if you use this dataset for your research.
 13"""
 14
 15import os
 16from tqdm import tqdm
 17from natsort import natsorted
 18from typing import Union, Tuple, List
 19
 20import json
 21import numpy as np
 22from skimage.draw import polygon
 23
 24import imageio.v3 as imageio
 25
 26from torch.utils.data import Dataset, DataLoader
 27
 28import torch_em
 29
 30from .. import util
 31
 32
 33URL = "https://drive.google.com/drive/folders/1TG9Bq-OaKkMXV2s42f_oZJdoTOfIAZLi"
 34
 35CATEGORIES = {0: "background", 1: "leiomyoma", 2: "GIST"}
 36"""The label ids of the lesion categories. The 'category_id' in the annotations (1: leiomyoma,
 372: GIST) is used directly as the pixel label, so that 0 marks background."""
 38
 39
 40def get_gist514_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 41    """Download the GIST514-DB dataset.
 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 downloaded.
 49    """
 50    data_dir = os.path.join(path, "usd514-db", "usd514_jpeg_roi")
 51    if os.path.exists(data_dir):
 52        return data_dir
 53
 54    os.makedirs(path, exist_ok=True)
 55    util.download_source_gdrive(path=path, url=URL, download=download, download_type="folder", expected_samples=1200)
 56
 57    return data_dir
 58
 59
 60def _rasterize_annotations(shape, image_annotations):
 61    labels = np.zeros(shape, dtype="uint8")
 62    for ann in image_annotations:
 63        category_id = ann["category_id"]
 64        for seg in ann["segmentation"]:
 65            c, r = np.asarray(seg[0::2]), np.asarray(seg[1::2])
 66            rr, cc = polygon(r, c, shape=shape)
 67            labels[rr, cc] = category_id
 68    return labels
 69
 70
 71def get_gist514_paths(path: Union[os.PathLike, str], download: bool = False) -> Tuple[List[str], List[str]]:
 72    """Get paths to the GIST514-DB data.
 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_gist514_data(path, download)
 83
 84    image_dir = os.path.join(data_dir, "images")
 85    annotation_path = os.path.join(data_dir, "annotations", "all_anno_crop.json")
 86    preprocessed_dir = os.path.join(data_dir, "preprocessed")
 87    os.makedirs(preprocessed_dir, exist_ok=True)
 88
 89    with open(annotation_path) as f:
 90        annotations = json.load(f)
 91
 92    image_paths, gt_paths = [], []
 93    for image_info in tqdm(annotations["images"], desc="Preprocessing labels"):
 94        image_path = os.path.join(image_dir, image_info["file_name"])
 95        if not os.path.exists(image_path):
 96            continue
 97
 98        fname = os.path.splitext(image_info["file_name"])[0]
 99        gt_path = os.path.join(preprocessed_dir, f"{fname}.tif")
100        if not os.path.exists(gt_path):
101            image_annotations = [a for a in annotations["annotations"] if a["image_id"] == image_info["id"]]
102            shape = (image_info["height"], image_info["width"])
103            labels = _rasterize_annotations(shape, image_annotations)
104            imageio.imwrite(gt_path, labels)
105
106        image_paths.append(image_path)
107        gt_paths.append(gt_path)
108
109    image_paths, gt_paths = natsorted(image_paths), natsorted(gt_paths)
110    return image_paths, gt_paths
111
112
113def get_gist514_dataset(
114    path: Union[os.PathLike, str],
115    patch_shape: Tuple[int, int],
116    resize_inputs: bool = False,
117    download: bool = False,
118    **kwargs
119) -> Dataset:
120    """Get the GIST514-DB dataset for lesion segmentation in endoscopic ultrasound images.
121
122    Args:
123        path: Filepath to a folder where the data is downloaded for further processing.
124        patch_shape: The patch shape to use for training.
125        resize_inputs: Whether to resize the inputs to the patch shape.
126        download: Whether to download the data if it is not present.
127        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
128
129    Returns:
130        The segmentation dataset.
131    """
132    image_paths, gt_paths = get_gist514_paths(path, download)
133
134    if resize_inputs:
135        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
136        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
137            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
138        )
139
140    return torch_em.default_segmentation_dataset(
141        raw_paths=image_paths,
142        raw_key=None,
143        label_paths=gt_paths,
144        label_key=None,
145        is_seg_dataset=False,
146        patch_shape=patch_shape,
147        **kwargs
148    )
149
150
151def get_gist514_loader(
152    path: Union[os.PathLike, str],
153    batch_size: int,
154    patch_shape: Tuple[int, int],
155    resize_inputs: bool = False,
156    download: bool = False,
157    **kwargs
158) -> DataLoader:
159    """Get the GIST514-DB dataloader for lesion segmentation in endoscopic ultrasound images.
160
161    Args:
162        path: Filepath to a folder where the data is downloaded for further processing.
163        batch_size: The batch size for training.
164        patch_shape: The patch shape to use for training.
165        resize_inputs: Whether to resize the inputs to the patch shape.
166        download: Whether to download the data if it is not present.
167        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
168
169    Returns:
170        The DataLoader.
171    """
172    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
173    dataset = get_gist514_dataset(path, patch_shape, resize_inputs, download, **ds_kwargs)
174    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)