torch_em.data.datasets.medical.platif

  1"""The PlaTiF dataset contains annotations for tibia bone segmentation in anteroposterior knee
  2radiographs, together with Schatzker classification labels for tibial plateau fractures.
  3
  4The dataset consists of 421 anteroposterior knee radiographs from 186 patients (normal and
  5fractured knees), collected at Shariati Hospital, Tehran University of Medical Sciences. Each
  6image comes with a manually produced (and MATLAB-refined) binary tibia segmentation mask, expert
  7validated, plus a Schatzker fracture type label (1-6) or a "no fracture" label (7).
  8
  9The dataset is located at https://doi.org/10.5281/zenodo.18007397 and is distributed under the
 10CC BY 4.0 license.
 11
 12This dataset is from the publication https://doi.org/10.1038/s41597-026-06560-5.
 13Please cite it if you use this dataset in your research.
 14"""
 15
 16import os
 17import shutil
 18from glob import glob
 19from tqdm import tqdm
 20from natsort import natsorted
 21from typing import Union, Tuple, List
 22
 23import imageio.v3 as imageio
 24
 25from torch.utils.data import Dataset, DataLoader
 26
 27import torch_em
 28
 29from .. import util
 30
 31
 32URLS = {
 33    "Patient Data_Part 1.zip": "https://zenodo.org/records/18007397/files/Patient%20Data_Part%201.zip",
 34    "Patient Data_Part 2.zip": "https://zenodo.org/records/18007397/files/Patient%20Data_Part%202.zip",
 35    "Patient Data_Part 3.zip": "https://zenodo.org/records/18007397/files/Patient%20Data_Part%203.zip",
 36    "Patient Data_Part 4.zip": "https://zenodo.org/records/18007397/files/Patient%20Data_Part%204.zip",
 37    "Patient Data_Part 5.zip": "https://zenodo.org/records/18007397/files/Patient%20Data_Part%205.zip",
 38}
 39
 40CHECKSUMS = {
 41    "Patient Data_Part 1.zip": "df72a6e0d988492eb87928105596001c6551db38d21dfb0a5d46942c45bd9d49",
 42    "Patient Data_Part 2.zip": "ef0f3315dbd9d47361dad23840cd4a8fbf36fbeb32955deddebc5fe29e14134b",
 43    "Patient Data_Part 3.zip": "63b77dc6e929df4d5c810a324d079905158d9196796ed9696ee5f1d81094ae09",
 44    "Patient Data_Part 4.zip": "d6135f67f5d7d76fab2b93ce059d56b6647e5bac87cf3b0d5ab0329121406756",
 45    "Patient Data_Part 5.zip": "dce5a29107a2e64eb7c25016e8cb65ba46d022f555983b144172602f2920416c",
 46}
 47
 48
 49def _preprocess_platif(mat_dir, preprocessed_dir):
 50    import scipy.io as sio
 51
 52    os.makedirs(os.path.join(preprocessed_dir, "images"), exist_ok=True)
 53    os.makedirs(os.path.join(preprocessed_dir, "masks"), exist_ok=True)
 54
 55    mat_paths = natsorted(glob(os.path.join(mat_dir, "**", "*.mat"), recursive=True))
 56    assert len(mat_paths) > 0, f"No '.mat' files were found at '{mat_dir}'."
 57
 58    for mat_path in tqdm(mat_paths, desc="Preprocessing inputs"):
 59        data = sio.loadmat(mat_path, simplify_cells=True)
 60        patient_keys = [k for k in data if not k.startswith("__")]
 61        assert len(patient_keys) == 1, f"Unexpected structure in '{mat_path}'."
 62        patient = data[patient_keys[0]]
 63
 64        view_keys = natsorted([k for k in patient if k.startswith("im")])
 65        for view_key in view_keys:
 66            view = patient[view_key]
 67
 68            image_path = os.path.join(preprocessed_dir, "images", f"{patient_keys[0]}_{view_key}.tif")
 69            mask_path = os.path.join(preprocessed_dir, "masks", f"{patient_keys[0]}_{view_key}.tif")
 70            if os.path.exists(image_path) and os.path.exists(mask_path):
 71                continue
 72
 73            image = (view["OriginalImage"] * 255).astype("uint8")
 74            mask = view["BW"].astype("uint8")
 75
 76            imageio.imwrite(image_path, image, compression="zlib")
 77            imageio.imwrite(mask_path, mask, compression="zlib")
 78
 79    shutil.rmtree(mat_dir)
 80
 81
 82def get_platif_data(path: Union[os.PathLike, str], download: bool = False) -> str:
 83    """Download the PlaTiF data.
 84
 85    Args:
 86        path: Filepath to a folder where the data is downloaded for further processing.
 87        download: Whether to download the data if it is not present.
 88
 89    Returns:
 90        Filepath where the preprocessed data is stored.
 91    """
 92    preprocessed_dir = os.path.join(path, "preprocessed")
 93    if os.path.exists(preprocessed_dir) and glob(os.path.join(preprocessed_dir, "images", "*.tif")):
 94        return preprocessed_dir
 95
 96    os.makedirs(path, exist_ok=True)
 97
 98    mat_dir = os.path.join(path, "mat_files")
 99    os.makedirs(mat_dir, exist_ok=True)
100    for fname, url in URLS.items():
101        zip_path = os.path.join(path, fname)
102        util.download_source(path=zip_path, url=url, download=download, checksum=CHECKSUMS[fname])
103        util.unzip(zip_path=zip_path, dst=mat_dir, remove=False)
104
105    _preprocess_platif(mat_dir, preprocessed_dir)
106    return preprocessed_dir
107
108
109def get_platif_paths(path: Union[os.PathLike, str], download: bool = False) -> Tuple[List[str], List[str]]:
110    """Get paths to the PlaTiF data.
111
112    Args:
113        path: Filepath to a folder where the data is downloaded for further processing.
114        download: Whether to download the data if it is not present.
115
116    Returns:
117        List of filepaths for the image data.
118        List of filepaths for the label data.
119    """
120    preprocessed_dir = get_platif_data(path, download)
121
122    image_paths = natsorted(glob(os.path.join(preprocessed_dir, "images", "*.tif")))
123    mask_paths = natsorted(glob(os.path.join(preprocessed_dir, "masks", "*.tif")))
124
125    if len(image_paths) == 0 or len(image_paths) != len(mask_paths):
126        raise RuntimeError("Something went wrong with fetching the image and label paths.")
127
128    return image_paths, mask_paths
129
130
131def get_platif_dataset(
132    path: Union[os.PathLike, str],
133    patch_shape: Tuple[int, int],
134    resize_inputs: bool = False,
135    download: bool = False,
136    **kwargs
137) -> Dataset:
138    """Get the PlaTiF dataset for tibia segmentation in knee radiographs.
139
140    Args:
141        path: Filepath to a folder where the data is downloaded for further processing.
142        patch_shape: The patch shape to use for training.
143        resize_inputs: Whether to resize inputs to the desired patch shape.
144        download: Whether to download the data if it is not present.
145        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
146
147    Returns:
148        The segmentation dataset.
149    """
150    image_paths, mask_paths = get_platif_paths(path, download)
151
152    if resize_inputs:
153        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
154        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
155            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
156        )
157
158    return torch_em.default_segmentation_dataset(
159        raw_paths=image_paths,
160        raw_key=None,
161        label_paths=mask_paths,
162        label_key=None,
163        patch_shape=patch_shape,
164        is_seg_dataset=False,
165        **kwargs
166    )
167
168
169def get_platif_loader(
170    path: Union[os.PathLike, str],
171    batch_size: int,
172    patch_shape: Tuple[int, int],
173    resize_inputs: bool = False,
174    download: bool = False,
175    **kwargs
176) -> DataLoader:
177    """Get the PlaTiF dataloader for tibia segmentation in knee radiographs.
178
179    Args:
180        path: Filepath to a folder where the data is downloaded for further processing.
181        batch_size: The batch size for training.
182        patch_shape: The patch shape to use for training.
183        resize_inputs: Whether to resize inputs to the desired patch shape.
184        download: Whether to download the data if it is not present.
185        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
186
187    Returns:
188        The DataLoader.
189    """
190    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
191    dataset = get_platif_dataset(path, patch_shape, resize_inputs, download, **ds_kwargs)
192    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)