torch_em.data.datasets.histopathology.hashi

  1"""This dataset contains annotations for invasive breast cancer region segmentation
  2in H&E-stained whole-slide histopathology tiles from four institutions.
  3
  4The data is from the publication https://doi.org/10.1371/journal.pone.0196828
  5("High-throughput adaptive sampling for whole-slide histopathology image analysis
  6(HASHI) via convolutional neural networks: application to invasive breast cancer
  7detection"). It is hosted on Zenodo at https://zenodo.org/records/4993672 under a
  8CC0-1.0 license. Please cite the publication if you use this dataset in your research.
  9
 10The dataset covers four cohorts: HUP and UHCMC/CWRU (used for training in the
 11publication), and CINJ and TCGA (held out for testing). Three pathologists manually
 12delineated invasive cancer regions at 2x magnification. Mask label values: 0
 13(background) and 1 (invasive cancer region).
 14
 15NOTE: Each cohort's images and masks are stored in separate multi-slide zip archives
 16on Zenodo. To avoid downloading a whole archive for a subset of its slides, each
 17raw/mask pair is fetched as a single zip member via an HTTP range request. Raw tiles
 18and their masks can differ by a 1px margin on each axis; each pair is cropped to their
 19shared shape on first download.
 20"""
 21
 22import os
 23import time
 24import struct
 25import zlib
 26import zipfile
 27from glob import glob
 28from natsort import natsorted
 29from typing import List, Literal, Optional, Tuple, Union
 30
 31from tqdm import tqdm
 32
 33import imageio.v3 as imageio
 34
 35import requests
 36
 37from torch.utils.data import Dataset, DataLoader
 38
 39import torch_em
 40
 41from .. import util
 42
 43
 44RECORD_URL = "https://zenodo.org/api/records/4993672/files"
 45N_RETRIES = 5
 46
 47COHORTS = {
 48    "hup": {"imgs": ["HUP_imgs_idx5_Part1.zip", "HUP_imgs_idx5_Part2.zip"], "masks": "HUP_masks.zip"},
 49    "cwru": {"imgs": ["CWRU_imgs_idx8.zip"], "masks": "CWRU_masks.zip"},
 50    "cinj": {"imgs": ["CINJ_imgs_idx5.zip"], "masks": "CINJ_masks_HG.zip"},
 51    "tcga": {"imgs": ["TCGA_imgs_idx5.zip"], "masks": "TCGA_masks.zip"},
 52}
 53
 54CHECKSUMS = {  # md5 of the full Zenodo archives, kept for provenance
 55    "HUP_imgs_idx5_Part1.zip": "c1bdaeb3c5bd2cd657b0081b9f52e3fe",
 56    "HUP_imgs_idx5_Part2.zip": "3127ee09f1752b76edd8d16383c52d30",
 57    "HUP_masks.zip": "6785a4cc69eae45217eb3e39843e3465",
 58    "CWRU_imgs_idx8.zip": "b1d13b2ecec81c0efe877bb25f45ece1",
 59    "CWRU_masks.zip": "a83e46a69d99c8e88ba5beef3be654a9",
 60    "CINJ_imgs_idx5.zip": "76dd6dc6f2e78bacdec427d0bd0d8740",
 61    "CINJ_masks_HG.zip": "65ab52631acaa40085972d1946ed1924",
 62    "TCGA_imgs_idx5.zip": "cc62330b4a2421219bd38f958e901c13",
 63    "TCGA_masks.zip": "a48c5c04c93d01831dc4a89a4dba609a",
 64}
 65
 66# HUP and CINJ raw tiles are named '{id}_idx5.png', with masks named differently.
 67# CWRU and TCGA raw tiles and masks share the exact same filename.
 68MASK_SUFFIX = {"hup": "_annotation_mask.png", "cinj": ".png"}
 69SPLITS = {"train": ["hup", "cwru"], "test": ["cinj", "tcga"]}
 70
 71
 72def _get_with_retries(url, **kwargs):
 73    # The Zenodo API gateway occasionally times out under range requests; retry with backoff.
 74    for attempt in range(N_RETRIES):
 75        try:
 76            r = requests.get(url, timeout=60, **kwargs)
 77            r.raise_for_status()
 78            return r
 79        except (requests.exceptions.RequestException,) as error:
 80            if attempt == N_RETRIES - 1:
 81                raise error
 82            time.sleep(2 ** attempt)
 83
 84
 85class _RemoteZipReader:
 86    """Seekable file-like object over a remote zip, for reading its (small) structural data."""
 87
 88    def __init__(self, url, size):
 89        self.url = url
 90        self.size = size
 91        self.pos = 0
 92
 93    def seek(self, offset, whence=0):
 94        if whence == 0:
 95            self.pos = offset
 96        elif whence == 1:
 97            self.pos += offset
 98        elif whence == 2:
 99            self.pos = self.size + offset
100        return self.pos
101
102    def tell(self):
103        return self.pos
104
105    def read(self, n=-1):
106        end = self.size - 1 if n is None or n < 0 else min(self.pos + n, self.size) - 1
107        if end < self.pos:
108            return b""
109        r = _get_with_retries(self.url, headers={"Range": f"bytes={self.pos}-{end}"})
110        data = r.content
111        self.pos += len(data)
112        return data
113
114    def readable(self):
115        return True
116
117    def seekable(self):
118        return True
119
120
121def _remote_size(url):
122    # A HEAD request fails on this host's presigned redirect target, which is only signed for GET.
123    r = _get_with_retries(url, stream=True, allow_redirects=True)
124    return int(r.headers["Content-Length"])
125
126
127def _list_zip_index(url):
128    """List a remote zip's central directory once, so individual members can be fetched by offset."""
129    size = _remote_size(url)
130    zf = zipfile.ZipFile(_RemoteZipReader(url, size))
131    return size, {info.filename: info for info in zf.infolist()}
132
133
134def _fetch_zip_member(url, size, info, dst_path):
135    if os.path.exists(dst_path):
136        return
137
138    # Fetch the local file header (with headroom for its variable-length fields) and the compressed
139    # payload in one range request. This dataset's members carry short names and no extra fields, so
140    # 256 bytes of headroom is always enough.
141    start = info.header_offset
142    end = start + 256 + info.compress_size - 1
143    r = _get_with_retries(url, headers={"Range": f"bytes={start}-{end}"})
144    buf = r.content
145
146    fn_len, extra_len = struct.unpack("<HH", buf[26:30])
147    data_start = 30 + fn_len + extra_len
148    compressed = buf[data_start:data_start + info.compress_size]
149    if len(compressed) != info.compress_size:
150        raise RuntimeError(f"Incomplete read for zip member '{info.filename}'.")
151
152    data = compressed if info.compress_type == zipfile.ZIP_STORED else zlib.decompressobj(-15).decompress(compressed)
153
154    os.makedirs(os.path.dirname(dst_path), exist_ok=True)
155    tmp_path = dst_path + ".tmp"
156    with open(tmp_path, "wb") as f:
157        f.write(data)
158    os.replace(tmp_path, dst_path)
159
160
161def _align_pair(img_path, mask_path):
162    """Crop a raw tile and its mask to their shared shape, in case they differ by a small margin."""
163    image = imageio.imread(img_path)
164    mask = imageio.imread(mask_path)
165    height, width = min(image.shape[0], mask.shape[0]), min(image.shape[1], mask.shape[1])
166    if image.shape[:2] != (height, width):
167        imageio.imwrite(img_path, image[:height, :width])
168    if mask.shape[:2] != (height, width):
169        imageio.imwrite(mask_path, mask[:height, :width])
170
171
172def _image_key(cohort, basename):
173    return basename[:-len("_idx5.png")] if cohort in MASK_SUFFIX else basename
174
175
176def _mask_key(cohort, basename):
177    suffix = MASK_SUFFIX.get(cohort)
178    return basename[:-len(suffix)] if suffix else basename
179
180
181def _resolve_cohorts(cohorts):
182    if cohorts is None:
183        return list(COHORTS)
184    if isinstance(cohorts, str):
185        cohorts = [cohorts]
186    invalid = set(cohorts) - set(COHORTS)
187    if invalid:
188        raise ValueError(f"Invalid cohort choices: {sorted(invalid)}. Choose from {sorted(COHORTS)}.")
189    return cohorts
190
191
192def get_hashi_data(
193    path: Union[os.PathLike, str],
194    cohorts: Optional[Union[str, List[str]]] = None,
195    sample_ids: Optional[List[str]] = None,
196    download: bool = False,
197) -> str:
198    """Download the HASHI invasive breast cancer segmentation data.
199
200    Args:
201        path: Filepath to a folder where the downloaded data will be saved.
202        cohorts: The cohort(s) to download. By default all four cohorts are downloaded.
203        sample_ids: The tile ids to restrict the data to. By default all tiles are used.
204        download: Whether to download the data if it is not present.
205
206    Returns:
207        Filepath to the folder where the raw images and masks are stored.
208    """
209    cohorts = _resolve_cohorts(cohorts)
210    path = str(path)
211    os.makedirs(path, exist_ok=True)
212
213    for cohort in cohorts:
214        masks_url = f"{RECORD_URL}/{COHORTS[cohort]['masks']}/content"
215        masks_size, masks_index = _list_zip_index(masks_url)
216        mask_by_key = {
217            _mask_key(cohort, name.rsplit("/", 1)[-1]): name
218            for name in masks_index if name.lower().endswith(".png")
219        }
220
221        for imgs_zip in COHORTS[cohort]["imgs"]:
222            imgs_url = f"{RECORD_URL}/{imgs_zip}/content"
223            imgs_size, imgs_index = _list_zip_index(imgs_url)
224            img_members = sorted(name for name in imgs_index if name.lower().endswith(".png"))
225
226            for img_name in tqdm(img_members, desc=f"Fetch {cohort} tiles ({imgs_zip})"):
227                basename = img_name.rsplit("/", 1)[-1]
228                key = _image_key(cohort, basename)
229                if sample_ids is not None and key not in sample_ids:
230                    continue
231                if key not in mask_by_key:
232                    raise RuntimeError(f"Missing mask for raw image '{img_name}' in cohort '{cohort}'.")
233                mask_name = mask_by_key[key]
234
235                img_path = os.path.join(path, cohort, "images", basename)
236                mask_path = os.path.join(path, cohort, "masks", mask_name.rsplit("/", 1)[-1])
237                if os.path.exists(img_path) and os.path.exists(mask_path):
238                    continue
239                if not download:
240                    raise RuntimeError(f"Data for cohort '{cohort}' is not found and download is set to False.")
241
242                _fetch_zip_member(imgs_url, imgs_size, imgs_index[img_name], img_path)
243                _fetch_zip_member(masks_url, masks_size, masks_index[mask_name], mask_path)
244                _align_pair(img_path, mask_path)
245
246    return path
247
248
249def get_hashi_paths(
250    path: Union[os.PathLike, str],
251    cohorts: Optional[Union[str, List[str]]] = None,
252    split: Optional[Literal["train", "test"]] = None,
253    sample_ids: Optional[List[str]] = None,
254    download: bool = False,
255) -> Tuple[List[str], List[str]]:
256    """Get paths to the HASHI invasive breast cancer segmentation images and masks.
257
258    Args:
259        path: Filepath to a folder where the downloaded data will be saved.
260        cohorts: The cohort(s) to use. By default all four cohorts are used.
261        split: The documented train ('hup', 'cwru') / test ('cinj', 'tcga') split. By default both are used.
262        sample_ids: The tile ids to restrict the data to. By default all tiles are used.
263        download: Whether to download the data if it is not present.
264
265    Returns:
266        List of filepaths for the image data.
267        List of filepaths for the label data.
268    """
269    cohorts = _resolve_cohorts(cohorts)
270    if split is not None:
271        cohorts = [cohort for cohort in cohorts if cohort in SPLITS[split]]
272
273    data_dir = get_hashi_data(path, cohorts, sample_ids, download)
274
275    raw_paths, label_paths = [], []
276    for cohort in cohorts:
277        img_paths = natsorted(glob(os.path.join(data_dir, cohort, "images", "*.png")))
278        mask_dir = os.path.join(data_dir, cohort, "masks")
279        mask_by_key = {_mask_key(cohort, name): name for name in os.listdir(mask_dir)}
280
281        for img_path in img_paths:
282            basename = os.path.basename(img_path)
283            key = _image_key(cohort, basename)
284            if sample_ids is not None and key not in sample_ids:
285                continue
286            if key not in mask_by_key:
287                raise RuntimeError(f"Missing mask for raw image '{img_path}' in cohort '{cohort}'.")
288            raw_paths.append(img_path)
289            label_paths.append(os.path.join(mask_dir, mask_by_key[key]))
290
291    if not raw_paths:
292        raise RuntimeError("Could not find any images and masks for the requested settings.")
293
294    return raw_paths, label_paths
295
296
297def get_hashi_dataset(
298    path: Union[os.PathLike, str],
299    patch_shape: Tuple[int, int],
300    cohorts: Optional[Union[str, List[str]]] = None,
301    split: Optional[Literal["train", "test"]] = None,
302    sample_ids: Optional[List[str]] = None,
303    resize_inputs: bool = False,
304    download: bool = False,
305    **kwargs,
306) -> Dataset:
307    """Get the HASHI dataset for invasive breast cancer region segmentation.
308
309    Args:
310        path: Filepath to a folder where the downloaded data will be saved.
311        patch_shape: The patch shape to use for training.
312        cohorts: The cohort(s) to use. By default all four cohorts are used.
313        split: The documented train ('hup', 'cwru') / test ('cinj', 'tcga') split. By default both are used.
314        sample_ids: The tile ids to restrict the data to. By default all tiles are used.
315        resize_inputs: Whether to resize the inputs.
316        download: Whether to download the data if it is not present.
317        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
318
319    Returns:
320        The segmentation dataset.
321    """
322    raw_paths, label_paths = get_hashi_paths(path, cohorts, split, sample_ids, download)
323
324    if resize_inputs:
325        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": True}
326        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
327            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
328        )
329
330    return torch_em.default_segmentation_dataset(
331        raw_paths=raw_paths,
332        raw_key=None,
333        label_paths=label_paths,
334        label_key=None,
335        patch_shape=patch_shape,
336        is_seg_dataset=False,
337        ndim=2,
338        with_channels=True,
339        **kwargs,
340    )
341
342
343def get_hashi_loader(
344    path: Union[os.PathLike, str],
345    batch_size: int,
346    patch_shape: Tuple[int, int],
347    cohorts: Optional[Union[str, List[str]]] = None,
348    split: Optional[Literal["train", "test"]] = None,
349    sample_ids: Optional[List[str]] = None,
350    resize_inputs: bool = False,
351    download: bool = False,
352    **kwargs,
353) -> DataLoader:
354    """Get the HASHI dataloader for invasive breast cancer region segmentation.
355
356    Args:
357        path: Filepath to a folder where the downloaded data will be saved.
358        batch_size: The batch size for training.
359        patch_shape: The patch shape to use for training.
360        cohorts: The cohort(s) to use. By default all four cohorts are used.
361        split: The documented train ('hup', 'cwru') / test ('cinj', 'tcga') split. By default both are used.
362        sample_ids: The tile ids to restrict the data to. By default all tiles are used.
363        resize_inputs: Whether to resize the inputs.
364        download: Whether to download the data if it is not present.
365        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or the PyTorch DataLoader.
366
367    Returns:
368        The DataLoader.
369    """
370    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
371    dataset = get_hashi_dataset(
372        path, patch_shape, cohorts, split, sample_ids, resize_inputs, download, **ds_kwargs
373    )
374    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)