torch_em.data.datasets.medical.hipas

  1"""The HiPaS dataset contains annotations for pulmonary artery and vein segmentation in non-contrast chest CT.
  2
  3The dataset consists of 250 non-contrast CT scans (512 x 512 in-plane, about 0.5 - 0.9 mm in-plane and 1 mm
  4slice spacing), each with a binary artery and a binary vein mask. `get_hipas_data` converts every case to a single
  5hdf5 file with the slice axis first (the keys are 'raw' and 'labels'), where the labels are a semantic volume
  6following `LABEL_IDS`: 1 = pulmonary artery, 2 = pulmonary vein. The artery and vein masks overlap in about 1 - 2 %
  7of their voxels, these voxels are labeled as vein.
  8
  9NOTE: This is not the same data as the already-integrated `torch_em.data.datasets.medical.airrc`, which provides
 10artery, vein and airway masks on resampled crops of the LUNA16 CT scans. HiPaS is a separate collection with its
 11own CT scans (the AirRC publication uses it as an independent external benchmark).
 12
 13The CT scans are stored in a single ~24 GB zip archive. To support downloading only a subset of the cases, the
 14individual scans are read from the archive with HTTP range requests instead of downloading it as a whole.
 15
 16The data is located at https://doi.org/10.5281/zenodo.14879605 (the example data of the publication, released under
 17an MIT license).
 18
 19This dataset is from the publication https://doi.org/10.1038/s41467-025-56505-6.
 20Please cite it if you use this dataset for your research.
 21"""
 22
 23import io
 24import os
 25import json
 26import uuid
 27import zlib
 28import struct
 29from glob import glob
 30from natsort import natsorted
 31from typing import Union, Tuple, List, Optional
 32from concurrent import futures
 33
 34import numpy as np
 35from tqdm import tqdm
 36
 37from torch.utils.data import Dataset, DataLoader
 38
 39import torch_em
 40
 41from .. import util
 42
 43
 44URL_BASE = "https://zenodo.org/api/records/14879605/files"
 45ANNOTATION_CHECKSUM = "8badcebd8c82d5ac0c1ea38d08d482efd37c78712009e7839f8339eee11f82ce"
 46
 47LABEL_IDS = {"artery": 1, "vein": 2}
 48"""Mapping from the name of a vessel type to its label id in the converted label volumes."""
 49
 50
 51def _read_zip_entries(url):
 52    import requests
 53
 54    size = int(requests.head(url, allow_redirects=True).headers["Content-Length"])
 55    tail = requests.get(url, headers={"Range": f"bytes={size - 65536}-{size - 1}"}).content
 56    eocd = tail.rfind(b"PK\x05\x06")
 57    n_entries, cd_size, cd_offset = struct.unpack("<HII", tail[eocd + 10:eocd + 20])
 58    if n_entries == 0xFFFF or cd_offset == 0xFFFFFFFF:
 59        locator = tail.rfind(b"PK\x06\x07")
 60        zip64_offset = struct.unpack("<Q", tail[locator + 8:locator + 16])[0]
 61        zip64 = requests.get(url, headers={"Range": f"bytes={zip64_offset}-{zip64_offset + 55}"}).content
 62        n_entries, cd_size, cd_offset = struct.unpack("<QQQ", zip64[32:56])
 63    directory = requests.get(url, headers={"Range": f"bytes={cd_offset}-{cd_offset + cd_size - 1}"}).content
 64
 65    entries, pos = {}, 0
 66    for _ in range(n_entries):
 67        fields = struct.unpack("<IHHHHHHIIIHHHHHII", directory[pos:pos + 46])
 68        method, csize, usize = fields[4], fields[8], fields[9]
 69        name_len, extra_len, comment_len = fields[10], fields[11], fields[12]
 70        header_offset = fields[16]
 71        name = directory[pos + 46:pos + 46 + name_len].decode()
 72        extra = directory[pos + 46 + name_len:pos + 46 + name_len + extra_len]
 73        offset = 0
 74        while offset < len(extra):
 75            tag, field_size = struct.unpack("<HH", extra[offset:offset + 4])
 76            if tag == 1:
 77                field, field_pos = extra[offset + 4:offset + 4 + field_size], 0
 78                if usize == 0xFFFFFFFF:
 79                    usize, field_pos = struct.unpack("<Q", field[field_pos:field_pos + 8])[0], field_pos + 8
 80                if csize == 0xFFFFFFFF:
 81                    csize, field_pos = struct.unpack("<Q", field[field_pos:field_pos + 8])[0], field_pos + 8
 82                if header_offset == 0xFFFFFFFF:
 83                    header_offset = struct.unpack("<Q", field[field_pos:field_pos + 8])[0]
 84            offset += 4 + field_size
 85        entries[name] = {"method": method, "compressed_size": csize, "header_offset": header_offset}
 86        pos += 46 + name_len + extra_len + comment_len
 87
 88    return entries
 89
 90
 91def _get_zip_entries(path, url):
 92    cache_path = os.path.join(path, "ct_scan_entries.json")
 93    if os.path.exists(cache_path):
 94        with open(cache_path) as f:
 95            return json.load(f)
 96
 97    entries = _read_zip_entries(url)
 98    tmp_path = f"{cache_path}.{uuid.uuid4().hex}.incomplete"
 99    with open(tmp_path, "w") as f:
100        json.dump(entries, f)
101    os.replace(tmp_path, cache_path)
102    return entries
103
104
105def _read_zip_member(url, entry):
106    import requests
107
108    offset = entry["header_offset"]
109    header = requests.get(url, headers={"Range": f"bytes={offset}-{offset + 29}"}).content
110    name_len, extra_len = struct.unpack("<HH", header[26:30])
111    start = offset + 30 + name_len + extra_len
112    response = requests.get(url, headers={"Range": f"bytes={start}-{start + entry['compressed_size'] - 1}"})
113    response.raise_for_status()
114    return zlib.decompress(response.content, -15) if entry["method"] == 8 else response.content
115
116
117def _convert_case(case_id, path, entries, ct_url):
118    import h5py
119
120    out_path = os.path.join(path, "preprocessed", f"{case_id}.h5")
121    if os.path.exists(out_path):
122        return
123
124    ct = np.load(io.BytesIO(_read_zip_member(ct_url, entries[f"ct_scan/{case_id}.npz"])))["data"]
125    labels = np.zeros(ct.shape, dtype="uint8")
126    for name, label_id in LABEL_IDS.items():
127        mask = np.load(os.path.join(path, "annotation", name, f"{case_id}.npz"))["data"]
128        assert mask.shape == ct.shape, f"{case_id}: {mask.shape} != {ct.shape}"
129        labels[mask > 0] = label_id
130
131    ct, labels = ct.transpose(2, 0, 1), labels.transpose(2, 0, 1)
132    chunks = (min(16, ct.shape[0]), 128, 128)
133    tmp_path = f"{out_path}.{uuid.uuid4().hex}.incomplete"
134    with h5py.File(tmp_path, "w") as f:
135        f.create_dataset("raw", data=ct, chunks=chunks, compression="gzip")
136        f.create_dataset("labels", data=labels, chunks=chunks, compression="gzip")
137    os.replace(tmp_path, out_path)
138
139
140def get_hipas_data(
141    path: Union[os.PathLike, str], n_cases: Optional[int] = None, n_workers: int = 4, download: bool = False
142) -> str:
143    """Download the HiPaS dataset and convert the cases to hdf5 files.
144
145    NOTE: The full collection is about 24 GB. Use `n_cases` to only download a subset for a quick start.
146
147    Args:
148        path: Filepath to a folder where the data is downloaded for further processing.
149        n_cases: The number of cases to download, sorted by case id. By default all 250 cases are downloaded.
150        n_workers: The number of parallel download and conversion workers.
151        download: Whether to download the data if it is not present.
152
153    Returns:
154        Filepath to the folder with the converted hdf5 files.
155    """
156    os.makedirs(path, exist_ok=True)
157
158    annotation_dir = os.path.join(path, "annotation")
159    if not os.path.exists(annotation_dir):
160        zip_path = os.path.join(path, "annotation.zip")
161        util.download_source(
162            path=zip_path, url=f"{URL_BASE}/annotation.zip/content", download=download, checksum=ANNOTATION_CHECKSUM,
163        )
164        util.unzip(zip_path=zip_path, dst=path, remove=False)
165
166    case_ids = natsorted(f[:-4] for f in os.listdir(os.path.join(annotation_dir, "artery")) if f.endswith(".npz"))
167    if n_cases is not None:
168        case_ids = case_ids[:n_cases]
169
170    preprocessed_dir = os.path.join(path, "preprocessed")
171    missing = [c for c in case_ids if not os.path.exists(os.path.join(preprocessed_dir, f"{c}.h5"))]
172    if not missing:
173        return preprocessed_dir
174    if not download:
175        raise RuntimeError(f"Cannot find the data at {path}, but download was set to False.")
176
177    os.makedirs(preprocessed_dir, exist_ok=True)
178    ct_url = f"{URL_BASE}/ct_scan.zip/content"
179    entries = _get_zip_entries(path, ct_url)
180    with futures.ThreadPoolExecutor(n_workers) as pool:
181        tasks = [pool.submit(_convert_case, case_id, path, entries, ct_url) for case_id in missing]
182        for task in tqdm(futures.as_completed(tasks), total=len(tasks), desc="Download and convert HiPaS"):
183            task.result()
184
185    return preprocessed_dir
186
187
188def get_hipas_paths(
189    path: Union[os.PathLike, str], n_cases: Optional[int] = None, download: bool = False,
190) -> List[str]:
191    """Get paths to the HiPaS data.
192
193    Args:
194        path: Filepath to a folder where the data is downloaded for further processing.
195        n_cases: The number of cases to use, sorted by case id. By default all 250 cases are used.
196        download: Whether to download the data if it is not present.
197
198    Returns:
199        List of filepaths for the hdf5 files, which contain the image data ('raw') and the label data ('labels').
200    """
201    data_dir = get_hipas_data(path, n_cases, download=download)
202    volume_paths = natsorted(glob(os.path.join(data_dir, "*.h5")))
203    if n_cases is not None:
204        volume_paths = volume_paths[:n_cases]
205    return volume_paths
206
207
208def get_hipas_dataset(
209    path: Union[os.PathLike, str],
210    patch_shape: Tuple[int, int, int],
211    n_cases: Optional[int] = None,
212    resize_inputs: bool = False,
213    download: bool = False,
214    **kwargs
215) -> Dataset:
216    """Get the HiPaS dataset for pulmonary artery and vein segmentation in non-contrast chest CT.
217
218    Args:
219        path: Filepath to a folder where the data is downloaded for further processing.
220        patch_shape: The patch shape to use for training.
221        n_cases: The number of cases to use, sorted by case id. By default all 250 cases are used.
222        resize_inputs: Whether to resize the inputs to the patch shape.
223        download: Whether to download the data if it is not present.
224        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
225
226    Returns:
227        The segmentation dataset.
228    """
229    volume_paths = get_hipas_paths(path, n_cases, download)
230
231    if resize_inputs:
232        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": False}
233        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
234            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
235        )
236
237    return torch_em.default_segmentation_dataset(
238        raw_paths=volume_paths,
239        raw_key="raw",
240        label_paths=volume_paths,
241        label_key="labels",
242        is_seg_dataset=True,
243        patch_shape=patch_shape,
244        ndim=3,
245        **kwargs
246    )
247
248
249def get_hipas_loader(
250    path: Union[os.PathLike, str],
251    batch_size: int,
252    patch_shape: Tuple[int, int, int],
253    n_cases: Optional[int] = None,
254    resize_inputs: bool = False,
255    download: bool = False,
256    **kwargs
257) -> DataLoader:
258    """Get the HiPaS dataloader for pulmonary artery and vein segmentation in non-contrast chest CT.
259
260    Args:
261        path: Filepath to a folder where the data is downloaded for further processing.
262        batch_size: The batch size for training.
263        patch_shape: The patch shape to use for training.
264        n_cases: The number of cases to use, sorted by case id. By default all 250 cases are used.
265        resize_inputs: Whether to resize the inputs to the patch shape.
266        download: Whether to download the data if it is not present.
267        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
268
269    Returns:
270        The DataLoader.
271    """
272    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
273    dataset = get_hipas_dataset(path, patch_shape, n_cases, resize_inputs, download, **ds_kwargs)
274    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)