torch_em.data.datasets.medical.hist_seg

  1"""The HIST-Seg dataset contains annotations for orthopedic tissue segmentation in hyperspectral surgical images.
  2
  3The dataset consists of 220 hyperspectral cubes of ten bovine specimens (22 per specimen), acquired with a HinaLea
  4camera in a simulated orthopedic surgery setting. Each cube has 299 spectral bands and 608 x 968 pixels and is
  5paired with a pixel-wise semantic mask with the classes given in `CLASS_NAMES`. NOTE: The label ids follow the
  6README of the dataset, so bone has the id 0 and there is no unlabeled id, the background of the operative field is
  7class 6.
  8
  9`get_hist_seg_data` converts every acquisition to a single hdf5 file. It contains the cube as 'raw' (bands first, shape
 10(299, 608, 968), float32 reflectance) and the class map as 'labels' (shape (608, 968)). The masks are distributed as
 11RGB images and are mapped to the class ids with the color table of the README.
 12
 13The cubes are stored in one ~15.5 GB zip archive per specimen (~155 GB in total). To support downloading only a
 14subset, single acquisitions are read from the archives with HTTP range requests (each member is verified with its
 15CRC32) instead of downloading the archives as a whole. Use `subjects` or `n_cases` to select a subset.
 16
 17NOTE: The dataset page sets a file access request flag, but the files are not restricted and can be downloaded
 18anonymously.
 19
 20The data is located at https://doi.org/10.57745/5N62WB, released under a CC-BY-4.0 license.
 21The code of the authors is available at https://github.com/lsllabisen/HIST-Seg.
 22Please cite the dataset if you use it for your research.
 23"""
 24
 25import os
 26import json
 27import uuid
 28import zlib
 29import struct
 30from io import BytesIO
 31from natsort import natsorted
 32from concurrent import futures
 33from typing import Union, Tuple, List, Optional, Sequence
 34
 35import numpy as np
 36from tqdm import tqdm
 37
 38from torch.utils.data import Dataset, DataLoader
 39
 40import torch_em
 41
 42from .. import util
 43
 44
 45DATAFILE_URL = "https://entrepot.recherche.data.gouv.fr/api/access/datafile/{}"
 46DATAFILE_IDS = {
 47    "B0001": 717111, "B0002": 717113, "B0003": 717132, "B0004": 717414, "B0005": 717496,
 48    "B0006": 717582, "B0007": 717597, "B0008": 717786, "B0009": 717820, "B0010": 717893,
 49}
 50SUBJECTS = list(DATAFILE_IDS)
 51
 52CLASS_NAMES = ["bone", "cartilage", "ligament", "flesh", "fat", "instruments", "background"]
 53"""The names of the classes. The label id of a class is its index in this list."""
 54
 55COLORS = [
 56    (142, 223, 124), (111, 225, 243), (245, 245, 63), (246, 58, 90), (243, 172, 99), (168, 109, 109), (169, 169, 172),
 57]
 58BLOCK_SIZE = 64 * 1024 * 1024
 59
 60
 61def _resolve(datafile_id):
 62    import requests
 63
 64    headers = {"User-Agent": "Mozilla/5.0", "Range": "bytes=0-0"}
 65    response = requests.get(DATAFILE_URL.format(datafile_id), headers=headers, allow_redirects=True)
 66    response.raise_for_status()
 67    return response.url, int(response.headers["Content-Range"].split("/")[1])
 68
 69
 70def _get_range(url, start, end):
 71    import requests
 72
 73    response = requests.get(url, headers={"Range": f"bytes={start}-{end}"})
 74    response.raise_for_status()
 75    return response.content
 76
 77
 78def _read_zip_entries(url, size):
 79    tail = _get_range(url, size - 65536, size - 1)
 80    eocd = tail.rfind(b"PK\x05\x06")
 81    n_entries, cd_size, cd_offset = struct.unpack("<HII", tail[eocd + 10:eocd + 20])
 82    if n_entries == 0xFFFF or cd_offset == 0xFFFFFFFF:
 83        locator = tail.rfind(b"PK\x06\x07")
 84        zip64_offset = struct.unpack("<Q", tail[locator + 8:locator + 16])[0]
 85        zip64 = _get_range(url, zip64_offset, zip64_offset + 55)
 86        n_entries, cd_size, cd_offset = struct.unpack("<QQQ", zip64[32:56])
 87    directory = _get_range(url, cd_offset, cd_offset + cd_size - 1)
 88
 89    entries, pos = {}, 0
 90    for _ in range(n_entries):
 91        fields = struct.unpack("<IHHHHHHIIIHHHHHII", directory[pos:pos + 46])
 92        method, crc, csize, usize = fields[4], fields[7], fields[8], fields[9]
 93        name_len, extra_len, comment_len = fields[10], fields[11], fields[12]
 94        header_offset = fields[16]
 95        name = directory[pos + 46:pos + 46 + name_len].decode()
 96        extra = directory[pos + 46 + name_len:pos + 46 + name_len + extra_len]
 97        offset = 0
 98        while offset < len(extra):
 99            tag, field_size = struct.unpack("<HH", extra[offset:offset + 4])
100            if tag == 1:
101                field, field_pos = extra[offset + 4:offset + 4 + field_size], 0
102                if usize == 0xFFFFFFFF:
103                    usize, field_pos = struct.unpack("<Q", field[field_pos:field_pos + 8])[0], field_pos + 8
104                if csize == 0xFFFFFFFF:
105                    csize, field_pos = struct.unpack("<Q", field[field_pos:field_pos + 8])[0], field_pos + 8
106                if header_offset == 0xFFFFFFFF:
107                    header_offset = struct.unpack("<Q", field[field_pos:field_pos + 8])[0]
108            offset += 4 + field_size
109        entries[name] = {"method": method, "crc": crc, "compressed_size": csize, "header_offset": header_offset}
110        pos += 46 + name_len + extra_len + comment_len
111
112    return entries
113
114
115def _get_entries(path, subject, download):
116    cache_path = os.path.join(path, f"entries_{subject}.json")
117    if os.path.exists(cache_path):
118        with open(cache_path) as f:
119            return json.load(f)
120
121    if not download:
122        raise RuntimeError(f"Cannot find the data at {path}, but download was set to False.")
123
124    url, size = _resolve(DATAFILE_IDS[subject])
125    entries = _read_zip_entries(url, size)
126    tmp_path = f"{cache_path}.{uuid.uuid4().hex}.incomplete"
127    with open(tmp_path, "w") as f:
128        json.dump(entries, f)
129    os.replace(tmp_path, cache_path)
130    return entries
131
132
133def _read_zip_member(url, entry):
134    offset = entry["header_offset"]
135    header = _get_range(url, offset, offset + 29)
136    name_len, extra_len = struct.unpack("<HH", header[26:30])
137    start = offset + 30 + name_len + extra_len
138
139    chunks = []
140    for block_start in range(0, entry["compressed_size"], BLOCK_SIZE):
141        block_end = min(block_start + BLOCK_SIZE, entry["compressed_size"]) - 1
142        chunks.append(_get_range(url, start + block_start, start + block_end))
143    data = b"".join(chunks)
144
145    if entry["method"] == 8:
146        data = zlib.decompress(data, -15)
147    if zlib.crc32(data) != entry["crc"]:
148        raise RuntimeError("The checksum of a downloaded archive member does not match.")
149    return data
150
151
152def _parse_header(header):
153    fields = {}
154    for line in header.decode(errors="replace").splitlines():
155        if "=" in line:
156            key, value = line.split("=", 1)
157            fields[key.strip()] = value.strip()
158    assert fields["interleave"] == "bsq" and fields["data type"] == "4" and fields["byte order"] == "0", fields
159    return int(fields["bands"]), int(fields["lines"]), int(fields["samples"])
160
161
162def _mask_to_labels(mask):
163    encoded = (mask[..., 0].astype("uint32") << 16) | (mask[..., 1].astype("uint32") << 8) | mask[..., 2]
164    labels = np.full(encoded.shape, 255, dtype="uint8")
165    for class_id, (r, g, b) in enumerate(COLORS):
166        labels[encoded == ((r << 16) | (g << 8) | b)] = class_id
167    assert (labels != 255).all(), "The mask contains a color that does not belong to a class."
168    return labels
169
170
171def _convert_acquisition(subject, acquisition, path, entries):
172    import h5py
173    from PIL import Image
174
175    out_path = os.path.join(path, "preprocessed", f"{acquisition}.h5")
176    if os.path.exists(out_path):
177        return
178
179    url, _ = _resolve(DATAFILE_IDS[subject])
180    prefix = f"{subject}/{acquisition}"
181    bands, lines, samples = _parse_header(_read_zip_member(url, entries[f"{prefix}/{acquisition}.hdr"]))
182    cube = np.frombuffer(_read_zip_member(url, entries[f"{prefix}/{acquisition}.dat"]), dtype="<f4")
183    cube = cube.reshape(bands, lines, samples)
184
185    mask = np.array(Image.open(BytesIO(_read_zip_member(url, entries[f"{prefix}/Annotations/{acquisition}.png"]))))
186    labels = _mask_to_labels(mask[..., :3])
187    assert labels.shape == cube.shape[1:], f"{acquisition}: {labels.shape} != {cube.shape[1:]}"
188
189    tmp_path = f"{out_path}.{uuid.uuid4().hex}.incomplete"
190    with h5py.File(tmp_path, "w") as f:
191        f.create_dataset("raw", data=cube, chunks=(bands, 64, 64), compression="gzip")
192        f.create_dataset("labels", data=labels, chunks=(64, 64), compression="gzip")
193    os.replace(tmp_path, out_path)
194
195
196def _list_acquisitions(path, subjects, n_cases, download):
197    acquisitions = []
198    for subject in subjects:
199        entries = _get_entries(path, subject, download)
200        names = natsorted({name.split("/")[1] for name in entries if name.endswith(".dat")})
201        acquisitions.extend(
202            (subject, name) for name in names if f"{subject}/{name}/Annotations/{name}.png" in entries
203        )
204        if n_cases is not None and len(acquisitions) >= n_cases:
205            break
206    return acquisitions[:n_cases]
207
208
209def get_hist_seg_data(
210    path: Union[os.PathLike, str],
211    subjects: Optional[Sequence[str]] = None,
212    n_cases: Optional[int] = None,
213    n_workers: int = 2,
214    download: bool = False,
215) -> str:
216    """Download the HIST-Seg dataset and convert the acquisitions to hdf5 files.
217
218    NOTE: The full collection is about 155 GB. Use `subjects` or `n_cases` to only download a subset.
219
220    Args:
221        path: Filepath to a folder where the data is downloaded for further processing.
222        subjects: The specimens to use, a subset of `SUBJECTS`. By default all ten specimens are used.
223        n_cases: The number of acquisitions to use, sorted by specimen and acquisition id. By default all are used.
224        n_workers: The number of parallel download and conversion workers. Each worker holds a cube of ~0.7 GB.
225        download: Whether to download the data if it is not present.
226
227    Returns:
228        Filepath to the folder with the converted hdf5 files.
229    """
230    subjects = SUBJECTS if subjects is None else list(subjects)
231    invalid = [s for s in subjects if s not in DATAFILE_IDS]
232    if invalid:
233        raise ValueError(f"{invalid} are not valid subjects. Choose from {SUBJECTS}.")
234
235    os.makedirs(path, exist_ok=True)
236    preprocessed_dir = os.path.join(path, "preprocessed")
237    acquisitions = _list_acquisitions(path, subjects, n_cases, download)
238
239    missing = [(s, a) for s, a in acquisitions if not os.path.exists(os.path.join(preprocessed_dir, f"{a}.h5"))]
240    if not missing:
241        return preprocessed_dir
242    if not download:
243        raise RuntimeError(f"Cannot find the data at {path}, but download was set to False.")
244
245    os.makedirs(preprocessed_dir, exist_ok=True)
246    entries = {subject: _get_entries(path, subject, download) for subject in {s for s, _ in missing}}
247    with futures.ThreadPoolExecutor(n_workers) as pool:
248        tasks = [pool.submit(_convert_acquisition, s, a, path, entries[s]) for s, a in missing]
249        for task in tqdm(futures.as_completed(tasks), total=len(tasks), desc="Download and convert HIST-Seg"):
250            task.result()
251
252    return preprocessed_dir
253
254
255def get_hist_seg_paths(
256    path: Union[os.PathLike, str],
257    subjects: Optional[Sequence[str]] = None,
258    n_cases: Optional[int] = None,
259    download: bool = False,
260) -> List[str]:
261    """Get paths to the HIST-Seg data.
262
263    Args:
264        path: Filepath to a folder where the data is downloaded for further processing.
265        subjects: The specimens to use, a subset of `SUBJECTS`. By default all ten specimens are used.
266        n_cases: The number of acquisitions to use, sorted by specimen and acquisition id. By default all are used.
267        download: Whether to download the data if it is not present.
268
269    Returns:
270        List of filepaths for the hdf5 files, which contain the cubes ('raw') and the class maps ('labels').
271    """
272    subjects = SUBJECTS if subjects is None else list(subjects)
273    preprocessed_dir = get_hist_seg_data(path, subjects, n_cases, download=download)
274    acquisitions = _list_acquisitions(path, subjects, n_cases, download)
275    return [os.path.join(preprocessed_dir, f"{a}.h5") for _, a in acquisitions]
276
277
278def get_hist_seg_dataset(
279    path: Union[os.PathLike, str],
280    patch_shape: Tuple[int, int],
281    subjects: Optional[Sequence[str]] = None,
282    n_cases: Optional[int] = None,
283    download: bool = False,
284    **kwargs
285) -> Dataset:
286    """Get the HIST-Seg dataset for tissue segmentation in hyperspectral surgical images.
287
288    The raw data has 299 channels (the spectral bands), the labels are the class ids of `CLASS_NAMES`.
289
290    Args:
291        path: Filepath to a folder where the data is downloaded for further processing.
292        patch_shape: The patch shape to use for training.
293        subjects: The specimens to use, a subset of `SUBJECTS`. By default all ten specimens are used.
294        n_cases: The number of acquisitions to use, sorted by specimen and acquisition id. By default all are used.
295        download: Whether to download the data if it is not present.
296        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
297
298    Returns:
299        The segmentation dataset.
300    """
301    volume_paths = get_hist_seg_paths(path, subjects, n_cases, download)
302
303    return torch_em.default_segmentation_dataset(
304        raw_paths=volume_paths,
305        raw_key="raw",
306        label_paths=volume_paths,
307        label_key="labels",
308        patch_shape=patch_shape,
309        is_seg_dataset=True,
310        with_channels=True,
311        ndim=2,
312        **kwargs
313    )
314
315
316def get_hist_seg_loader(
317    path: Union[os.PathLike, str],
318    batch_size: int,
319    patch_shape: Tuple[int, int],
320    subjects: Optional[Sequence[str]] = None,
321    n_cases: Optional[int] = None,
322    download: bool = False,
323    **kwargs
324) -> DataLoader:
325    """Get the HIST-Seg dataloader for tissue segmentation in hyperspectral surgical images.
326
327    Args:
328        path: Filepath to a folder where the data is downloaded for further processing.
329        batch_size: The batch size for training.
330        patch_shape: The patch shape to use for training.
331        subjects: The specimens to use, a subset of `SUBJECTS`. By default all ten specimens are used.
332        n_cases: The number of acquisitions to use, sorted by specimen and acquisition id. By default all are used.
333        download: Whether to download the data if it is not present.
334        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
335
336    Returns:
337        The DataLoader.
338    """
339    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
340    dataset = get_hist_seg_dataset(path, patch_shape, subjects, n_cases, download, **ds_kwargs)
341    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)