torch_em.data.datasets.histopathology.peso

  1"""This dataset contains annotations for prostate epithelium segmentation in
  2H&E-stained whole-slide histopathology images of prostatectomy specimens.
  3
  4The data is from the publication https://doi.org/10.1038/s41598-018-37257-4
  5("Epithelium segmentation using deep learning in H&E-stained prostate specimens
  6with immunohistochemistry as reference standard"). It is hosted on Zenodo at
  7https://zenodo.org/records/1485967 under a CC BY-NC-SA 4.0 license.
  8Please cite the publication if you use this dataset.
  9
 10This loader exposes the 25 training whole-slide images whose epithelium mask
 11received manual pathologist correction. The remaining 37 training masks are
 12uncorrected predictions of an IHC-guided U-Net, and the 40 test slides ship
 13only region outlines and a benign/cancer label rather than a raster mask, so
 14neither is included here.
 15
 16Mask label values: 0 (unannotated), 1 (non-epithelium tissue within an
 17annotated region), and 2 (epithelium). Only small regions of interest within
 18each slide are annotated (about 7% of a slide's area); the rest of a slide's
 19mask is unannotated, so `get_peso_dataset` / `get_peso_loader` default to a
 20`MinForegroundSampler` that rejects mostly-unannotated patches.
 21
 22NOTE: The whole-slide images and masks are multi-resolution pyramidal TIFFs of
 23several gigabytes each, bundled inside a few large multi-slide zip archives on
 24Zenodo. To avoid downloading unrelated slides, each requested slide is fetched
 25as a single zip member via an HTTP range request rather than the whole archive.
 26On first use each requested slide is also converted into a chunked HDF5 file at
 27the requested pyramid level, which requires some time and disk space.
 28"""
 29
 30import os
 31import struct
 32import zipfile
 33import zlib
 34from pathlib import Path
 35from typing import List, Optional, Tuple, Union
 36
 37from tqdm import tqdm
 38
 39import requests
 40
 41import torch
 42from torch.utils.data import Dataset, DataLoader
 43
 44import torch_em
 45from torch_em.data.sampler import MinForegroundSampler
 46
 47from .. import util
 48
 49
 50RECORD_URL = "https://zenodo.org/api/records/1485967/files"
 51
 52CHECKSUMS = {  # md5 of the full Zenodo archives, kept for provenance
 53    "peso_training_masks_corrected.zip": "8e2c86fcecfafe09c9d48a60b42441b5",
 54    "peso_training_wsi_1.zip": "8e4e53d7ba855fc2f318dce94b05fe31",
 55    "peso_training_wsi_2.zip": "bfcae8b444c12c0ecbb717dc37334020",
 56    "peso_training_wsi_3.zip": "f7d484acec429c3a9d9e685969edd82b",
 57    "peso_training_wsi_4.zip": "75a350b450193b48e9bed3a3484f639d",
 58    "peso_training_wsi_5.zip": "bc94373db95c5e2ceefe08c45a8e54db",
 59    "peso_training_wsi_6.zip": "4fa8cdd4b748d67b6c39a982be7b627a",
 60}
 61
 62# The 25 corrected slides, mapped to the wsi archive that holds their raw image.
 63STEM_TO_WSI_ZIP = {
 64    "pds_6": 1, "pds_8": 1,
 65    "pds_34": 2,
 66    "pds_35": 3, "pds_38": 3, "pds_39": 3, "pds_40": 3, "pds_43": 3,
 67    "pds_46": 4, "pds_56": 4, "pds_60": 4, "pds_64": 4,
 68    "pds_69": 5, "pds_70": 5, "pds_71": 5, "pds_72": 5, "pds_73": 5, "pds_79": 5,
 69    "pds_91": 6, "pds_93": 6, "pds_96": 6, "pds_99": 6, "pds_100": 6, "pds_101": 6, "pds_102": 6,
 70}
 71
 72
 73class _RemoteZipReader:
 74    """Seekable file-like object over a remote zip, for reading its (small) structural data."""
 75
 76    def __init__(self, url, size):
 77        self.url = url
 78        self.size = size
 79        self.pos = 0
 80
 81    def seek(self, offset, whence=0):
 82        if whence == 0:
 83            self.pos = offset
 84        elif whence == 1:
 85            self.pos += offset
 86        elif whence == 2:
 87            self.pos = self.size + offset
 88        return self.pos
 89
 90    def tell(self):
 91        return self.pos
 92
 93    def read(self, n=-1):
 94        end = self.size - 1 if n is None or n < 0 else min(self.pos + n, self.size) - 1
 95        if end < self.pos:
 96            return b""
 97        r = requests.get(self.url, headers={"Range": f"bytes={self.pos}-{end}"})
 98        r.raise_for_status()
 99        data = r.content
100        self.pos += len(data)
101        return data
102
103    def readable(self):
104        return True
105
106    def seekable(self):
107        return True
108
109
110def _download_zip_member(archive_url, member, dst_path, chunk_size=1024 * 1024 * 8):
111    """Fetch one member of a large remote zip via a single HTTP range request for its
112    compressed bytes, instead of downloading the whole multi-slide archive.
113    """
114    if os.path.exists(dst_path):
115        return
116
117    size = int(requests.head(archive_url, allow_redirects=True).headers["Content-Length"])
118    zf = zipfile.ZipFile(_RemoteZipReader(archive_url, size))
119    info = zf.getinfo(member)
120
121    header_reader = _RemoteZipReader(archive_url, size)
122    header_reader.seek(info.header_offset)
123    local_header = header_reader.read(30)
124    fn_len, extra_len = struct.unpack("<HH", local_header[26:30])
125    data_start = info.header_offset + 30 + fn_len + extra_len
126    data_end = data_start + info.compress_size - 1
127
128    part_path = dst_path + ".part"
129    with requests.get(archive_url, headers={"Range": f"bytes={data_start}-{data_end}"}, stream=True) as r:
130        r.raise_for_status()
131        with open(part_path, "wb") as f:
132            for chunk in r.iter_content(chunk_size=chunk_size):
133                f.write(chunk)
134
135    if info.compress_type == zipfile.ZIP_STORED:
136        os.replace(part_path, dst_path)
137        return
138
139    decompressor = zlib.decompressobj(-15)
140    with open(part_path, "rb") as src, open(dst_path, "wb") as dst:
141        while True:
142            chunk = src.read(chunk_size)
143            if not chunk:
144                break
145            dst.write(decompressor.decompress(chunk))
146        dst.write(decompressor.flush())
147    os.remove(part_path)
148
149
150def _resolve_sample_ids(sample_ids):
151    stems = sorted(STEM_TO_WSI_ZIP, key=lambda s: int(s.split("_")[1]))
152    if sample_ids is None:
153        return stems
154    missing = sorted(set(sample_ids) - set(stems))
155    if missing:
156        raise ValueError(f"The following sample ids are not part of this dataset: {missing}")
157    return [stem for stem in stems if stem in sample_ids]
158
159
160def _open_level(series, level_index):
161    import zarr
162
163    # The pyramidal TIFFs are natively tiled, so a zarr view reads only the requested tiles lazily.
164    array = zarr.open(series.aszarr(), mode="r")
165    return array if hasattr(array, "shape") else array[str(level_index)]
166
167
168def _convert_slide(image_path, mask_path, output_path, resolution_level, tile=4096):
169    import h5py
170    import tifffile
171
172    image_series = tifffile.TiffFile(image_path).series[0]
173    mask_series = tifffile.TiffFile(mask_path).series[0]
174
175    # The released masks were rasterized one pyramid level finer than the raw slide.
176    mask_level = resolution_level + 1
177    height, width = image_series.levels[resolution_level].shape[:2]
178    mask_height, mask_width = mask_series.levels[mask_level].shape[:2]
179    if abs(height - mask_height) > 1 or abs(width - mask_width) > 1:
180        raise RuntimeError(
181            f"The mask '{mask_path}' does not match the raw shape ({height}, {width}) "
182            f"at level {resolution_level}: got ({mask_height}, {mask_width}) at mask level {mask_level}."
183        )
184    height, width = min(height, mask_height), min(width, mask_width)
185
186    image = _open_level(image_series, resolution_level)
187    mask = _open_level(mask_series, mask_level)
188
189    tmp_path = output_path + ".tmp"
190    with h5py.File(tmp_path, "w") as f:
191        raw = f.create_dataset(
192            "images/raw", shape=(3, height, width), dtype="uint8", compression="gzip", chunks=(1, 512, 512)
193        )
194        labels = f.create_dataset(
195            "labels/mask", shape=(height, width), dtype="uint8", compression="gzip", chunks=(512, 512)
196        )
197        for y in tqdm(range(0, height, tile), desc=f"Converting {Path(image_path).stem}"):
198            for x in range(0, width, tile):
199                th, tw = min(tile, height - y), min(tile, width - x)
200                raw[:, y:y + th, x:x + tw] = image[y:y + th, x:x + tw].transpose(2, 0, 1)
201                labels[y:y + th, x:x + tw] = mask[y:y + th, x:x + tw]
202
203    os.replace(tmp_path, output_path)
204
205
206def get_peso_data(
207    path: Union[os.PathLike, str],
208    sample_ids: Optional[List[str]] = None,
209    resolution_level: int = 0,
210    download: bool = False,
211) -> str:
212    """Download and preprocess the PESO prostate epithelium segmentation data.
213
214    Args:
215        path: Filepath to a folder where the data will be saved.
216        sample_ids: The slide stems to restrict the data to, e.g. ['pds_8', 'pds_34'].
217            By default all 25 corrected slides are used.
218        resolution_level: The pyramid level to convert. 0 is the native raw resolution;
219            use a higher level to reduce the size of the preprocessed data.
220        download: Whether to download the data if it is not present.
221
222    Returns:
223        Filepath to the folder where the preprocessed data is stored.
224    """
225    stems = _resolve_sample_ids(sample_ids)
226
227    raw_dir = os.path.join(path, "raw")
228    mask_dir = os.path.join(path, "masks")
229    preprocessed_dir = os.path.join(path, "preprocessed")
230    os.makedirs(preprocessed_dir, exist_ok=True)
231
232    for stem in stems:
233        output_path = os.path.join(preprocessed_dir, f"{stem}_level{resolution_level}.h5")
234        if os.path.exists(output_path):
235            continue
236
237        os.makedirs(raw_dir, exist_ok=True)
238        os.makedirs(mask_dir, exist_ok=True)
239        image_path = os.path.join(raw_dir, f"{stem}_HE.tif")
240        mask_path = os.path.join(mask_dir, f"{stem}_HE_training_mask_corrected.tif")
241
242        if not (os.path.exists(image_path) and os.path.exists(mask_path)):
243            if not download:
244                raise RuntimeError(f"Data for '{stem}' is not found and download is set to False.")
245
246            wsi_url = f"{RECORD_URL}/peso_training_wsi_{STEM_TO_WSI_ZIP[stem]}.zip/content"
247            masks_url = f"{RECORD_URL}/peso_training_masks_corrected.zip/content"
248            _download_zip_member(wsi_url, f"{stem}_HE.tif", image_path)
249            _download_zip_member(masks_url, f"{stem}_HE_training_mask_corrected.tif", mask_path)
250
251        _convert_slide(image_path, mask_path, output_path, resolution_level)
252
253    return preprocessed_dir
254
255
256def get_peso_paths(
257    path: Union[os.PathLike, str],
258    sample_ids: Optional[List[str]] = None,
259    resolution_level: int = 0,
260    download: bool = False,
261) -> List[str]:
262    """Get paths to the PESO prostate epithelium segmentation data.
263
264    Args:
265        path: Filepath to a folder where the data will be saved.
266        sample_ids: The slide stems to restrict the data to, e.g. ['pds_8', 'pds_34'].
267            By default all 25 corrected slides are used.
268        resolution_level: The pyramid level to convert. 0 is the native raw resolution;
269            use a higher level to reduce the size of the preprocessed data.
270        download: Whether to download the data if it is not present.
271
272    Returns:
273        List of filepaths to the preprocessed HDF5 files.
274    """
275    preprocessed_dir = get_peso_data(path, sample_ids, resolution_level, download)
276    stems = _resolve_sample_ids(sample_ids)
277    return [os.path.join(preprocessed_dir, f"{stem}_level{resolution_level}.h5") for stem in stems]
278
279
280def get_peso_dataset(
281    path: Union[os.PathLike, str],
282    patch_shape: Tuple[int, int],
283    sample_ids: Optional[List[str]] = None,
284    resolution_level: int = 0,
285    download: bool = False,
286    label_dtype: torch.dtype = torch.int64,
287    resize_inputs: bool = False,
288    **kwargs
289) -> Dataset:
290    """Get the PESO dataset for prostate epithelium segmentation in whole-slide histopathology images.
291
292    Args:
293        path: Filepath to a folder where the data will be saved.
294        patch_shape: The patch shape to use for training.
295        sample_ids: The slide stems to restrict the data to, e.g. ['pds_8', 'pds_34'].
296            By default all 25 corrected slides are used.
297        resolution_level: The pyramid level to convert. 0 is the native raw resolution;
298            use a higher level to reduce the size of the preprocessed data.
299        download: Whether to download the data if it is not present.
300        label_dtype: The datatype of the labels.
301        resize_inputs: Whether to resize the input images.
302        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
303
304    Returns:
305        The segmentation dataset.
306    """
307    volume_paths = get_peso_paths(path, sample_ids, resolution_level, download)
308
309    if resize_inputs:
310        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": True}
311        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
312            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
313        )
314
315    # Only small regions of interest are annotated per slide, so most patches would otherwise be unannotated.
316    kwargs.setdefault("sampler", MinForegroundSampler(min_fraction=0.05, background_id=0, p_reject=0.9))
317
318    return torch_em.default_segmentation_dataset(
319        raw_paths=volume_paths,
320        raw_key="images/raw",
321        label_paths=volume_paths,
322        label_key="labels/mask",
323        patch_shape=patch_shape,
324        label_dtype=label_dtype,
325        is_seg_dataset=True,
326        with_channels=True,
327        ndim=2,
328        **kwargs
329    )
330
331
332def get_peso_loader(
333    path: Union[os.PathLike, str],
334    patch_shape: Tuple[int, int],
335    batch_size: int,
336    sample_ids: Optional[List[str]] = None,
337    resolution_level: int = 0,
338    download: bool = False,
339    label_dtype: torch.dtype = torch.int64,
340    resize_inputs: bool = False,
341    **kwargs
342) -> DataLoader:
343    """Get the PESO dataloader for prostate epithelium segmentation in whole-slide histopathology images.
344
345    Args:
346        path: Filepath to a folder where the data will be saved.
347        patch_shape: The patch shape to use for training.
348        batch_size: The batch size for training.
349        sample_ids: The slide stems to restrict the data to, e.g. ['pds_8', 'pds_34'].
350            By default all 25 corrected slides are used.
351        resolution_level: The pyramid level to convert. 0 is the native raw resolution;
352            use a higher level to reduce the size of the preprocessed data.
353        download: Whether to download the data if it is not present.
354        label_dtype: The datatype of the labels.
355        resize_inputs: Whether to resize the input images.
356        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or the PyTorch DataLoader.
357
358    Returns:
359        The DataLoader.
360    """
361    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
362    dataset = get_peso_dataset(
363        path=path, patch_shape=patch_shape, sample_ids=sample_ids, resolution_level=resolution_level,
364        download=download, label_dtype=label_dtype, resize_inputs=resize_inputs, **ds_kwargs
365    )
366    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)