torch_em.data.datasets.medical.btxrd

  1"""BTXRD is a dataset for the segmentation of primary bone tumors in radiographs.
  2
  3This dataset is located at https://doi.org/10.6084/m9.figshare.27865398.
  4The dataset is from the publication https://doi.org/10.1038/s41597-024-04311-y.
  5Please cite it if you use this dataset for your research.
  6"""
  7
  8import os
  9import json
 10from glob import glob
 11from tqdm import tqdm
 12from pathlib import Path
 13from natsort import natsorted
 14from typing import Union, Tuple, List
 15
 16import numpy as np
 17from skimage import draw
 18import imageio.v3 as imageio
 19
 20import torch_em
 21
 22from .. import util
 23from ..light_microscopy.neurips_cell_seg import to_rgb
 24
 25
 26URL = "https://ndownloader.figshare.com/files/50653575"
 27CHECKSUM = "e7c800c3b4e090262b160525a0765f9d93bcc53d639c03806a3ac47b0ded3373"
 28
 29
 30def get_btxrd_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 31    """Download the BTXRD 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, "BTXRD")
 41    if os.path.exists(data_dir):
 42        return data_dir
 43
 44    os.makedirs(path, exist_ok=True)
 45
 46    zip_path = os.path.join(path, "BTXRD.zip")
 47    util.download_source(path=zip_path, url=URL, download=download, checksum=CHECKSUM)
 48    util.unzip(zip_path=zip_path, dst=path)
 49
 50    return data_dir
 51
 52
 53def _create_mask(annotation_path, image_shape):
 54    with open(annotation_path) as f:
 55        annotation = json.load(f)
 56
 57    mask = np.zeros(image_shape[:2], dtype=np.uint8)
 58    for shape in annotation["shapes"]:
 59        if shape["shape_type"] != "polygon":
 60            continue
 61
 62        points = np.array(shape["points"])
 63        rr, cc = draw.polygon(points[:, 1], points[:, 0], shape=mask.shape)
 64        mask[rr, cc] = 1
 65
 66    return mask
 67
 68
 69def _preprocess_labels(data_dir):
 70    image_dir = os.path.join(data_dir, "images")
 71    annotation_dir = os.path.join(data_dir, "Annotations")
 72    gt_dir = os.path.join(data_dir, "masks")
 73    os.makedirs(gt_dir, exist_ok=True)
 74
 75    annotation_paths = natsorted(glob(os.path.join(annotation_dir, "*.json")))
 76
 77    image_paths, gt_paths = [], []
 78    for annotation_path in tqdm(annotation_paths, desc="Converting polygon annotations to masks"):
 79        image_id = Path(annotation_path).stem
 80
 81        image_path = glob(os.path.join(image_dir, f"{image_id}.jp*g"))
 82        assert len(image_path) == 1, f"Could not find a unique matching image for '{image_id}'."
 83        image_path = image_path[0]
 84
 85        gt_path = os.path.join(gt_dir, f"{image_id}.tif")
 86        if not os.path.exists(gt_path):
 87            image_shape = imageio.imread(image_path).shape
 88            mask = _create_mask(annotation_path, image_shape)
 89            imageio.imwrite(gt_path, mask)
 90
 91        image_paths.append(image_path)
 92        gt_paths.append(gt_path)
 93
 94    return image_paths, gt_paths
 95
 96
 97def get_btxrd_paths(path: Union[os.PathLike, str], download: bool = False) -> Tuple[List[str], List[str]]:
 98    """Get paths to the BTXRD data.
 99
100    Args:
101        path: Filepath to a folder where the data is downloaded for further processing.
102        download: Whether to download the data if it is not present.
103
104    Returns:
105        List of filepaths for the image data.
106        List of filepaths for the label data.
107    """
108    data_dir = get_btxrd_data(path=path, download=download)
109    image_paths, gt_paths = _preprocess_labels(data_dir)
110    return image_paths, gt_paths
111
112
113def get_btxrd_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):
120    """Get the BTXRD dataset for segmentation of primary bone tumors in radiographs.
121
122    Args:
123        path: Filepath to a folder where the downloaded data will be saved.
124        patch_shape: The patch shape to use for training.
125        resize_inputs: Whether to resize the inputs to the expected 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_btxrd_paths(path=path, download=download)
133
134    if resize_inputs:
135        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": True}
136        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
137            kwargs=kwargs,
138            patch_shape=patch_shape,
139            resize_inputs=resize_inputs,
140            resize_kwargs=resize_kwargs,
141            ensure_rgb=to_rgb,
142        )
143
144    return torch_em.default_segmentation_dataset(
145        raw_paths=image_paths,
146        raw_key=None,
147        label_paths=gt_paths,
148        label_key=None,
149        patch_shape=patch_shape,
150        is_seg_dataset=False,
151        **kwargs
152    )
153
154
155def get_btxrd_loader(
156    path: Union[os.PathLike, str],
157    batch_size: int,
158    patch_shape: Tuple[int, int],
159    resize_inputs: bool = False,
160    download: bool = False,
161    **kwargs
162):
163    """Get the BTXRD dataloader for segmentation of primary bone tumors in radiographs.
164
165    Args:
166        path: Filepath to a folder where the downloaded data will be saved.
167        batch_size: The batch size for training.
168        patch_shape: The patch shape to use for training.
169        resize_inputs: Whether to resize the inputs to the expected patch shape.
170        download: Whether to download the data if it is not present.
171        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
172
173    Returns:
174        The DataLoader.
175    """
176    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
177    dataset = get_btxrd_dataset(path, patch_shape, resize_inputs, download, **ds_kwargs)
178    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)