torch_em.data.datasets.histopathology.crc_epithelium

  1"""This dataset contains annotations for epithelium segmentation in colorectal tissue microarray
  2cores, stained with H&E, 13 immunohistochemistry (IHC) protein markers, and in situ hybridization
  3(ISH) for two microRNAs plus positive/negative controls.
  4
  5The data is from the publication https://doi.org/10.1111/apm.70051 ("miR-143 and miR-145 in
  6Colorectal Cancer: A Digital Pathology Approach on Expressions and Protein Correlations"). It is
  7hosted on DataverseNO at https://doi.org/10.18710/DIGQGQ under a CC0-1.0 license. Please cite the
  8publication if you use this dataset in your research.
  9
 10The dataset covers 100 patients, each with three normal mucosa and three cancer tissue microarray
 11cores per stain. Mask label values: 0 (background) and 1 (epithelium).
 12
 13NOTE: Each stain is stored on DataverseNO as one multi-gigabyte zip archive of all its cores. To
 14avoid downloading a whole archive for the images actually requested, each raw/mask pair is fetched
 15as a single zip member via an HTTP range request.
 16"""
 17
 18import os
 19import struct
 20import zlib
 21import zipfile
 22from glob import glob
 23from pathlib import Path
 24from typing import List, Literal, Optional, Tuple, Union
 25
 26from tqdm import tqdm
 27
 28import requests
 29
 30from torch.utils.data import Dataset, DataLoader
 31
 32import torch_em
 33
 34from .. import util
 35
 36
 37BASE_URL = "https://dataverse.no/api/access/datafile"
 38
 39URLS = {
 40    "HE": f"{BASE_URL}/233917",
 41    "ECAD": f"{BASE_URL}/233915",
 42    "VIMENTIN": f"{BASE_URL}/233933",
 43    "SMA": f"{BASE_URL}/233929",
 44    "Ki67": f"{BASE_URL}/233920",
 45    "SMAD3": f"{BASE_URL}/233934",
 46    "MACC1": f"{BASE_URL}/233918",
 47    "LASP1": f"{BASE_URL}/233921",
 48    "CD44": f"{BASE_URL}/233913",
 49    "NAIP": f"{BASE_URL}/233928",
 50    "KLF5": f"{BASE_URL}/233919",
 51    "FSCN1": f"{BASE_URL}/233916",
 52    "CTNND1": f"{BASE_URL}/233914",
 53    "KRAS": f"{BASE_URL}/233922",
 54    "miR-143": f"{BASE_URL}/234318",
 55    "miR-145": f"{BASE_URL}/233924",
 56    "U6": f"{BASE_URL}/233932",
 57    "Scr": f"{BASE_URL}/233930",
 58}
 59
 60CHECKSUMS = {  # md5 of the full DataverseNO archives, kept for provenance
 61    "HE": "f82f503e821ee4c546ce3e5b82465cb7",
 62    "ECAD": "b83f719ebb1e46dbbfdfb17a784e467f",
 63    "VIMENTIN": "842bfcac56034fcb264d81c88e88455f",
 64    "SMA": "0debe0bf6613b6d5651c45fb7b84581d",
 65    "Ki67": "4060c46fa1334c3fe804346e26bf9225",
 66    "SMAD3": "f8a3cc5e556c833a384faa5914f66f08",
 67    "MACC1": "97b223362abcfaf3db010794300cf470",
 68    "LASP1": "d116d8cf53532a9729d25dbf5ee11ea1",
 69    "CD44": "c05e00ddb394f2d82cff3039c9fd667d",
 70    "NAIP": "b83bdbcd9211b94d1eba86353347e142",
 71    "KLF5": "fb48fb9dccb9de02b3c0e43f3f348410",
 72    "FSCN1": "1acfcc09f439eee9da58b26c836d980b",
 73    "CTNND1": "0866a21e48de86b7ff01daef02a7cea2",
 74    "KRAS": "63cee667a452c6f1e236b719d2eb78d3",
 75    "miR-143": "0bda5e92b2b4b3cf269879fb5f7e48ab",
 76    "miR-145": "76ee831cd078245ea54d78405b5dd8b4",
 77    "U6": "14918f009194e342843d284bd4a4d380",
 78    "Scr": "1e401fde4d37a721944cc23fc6b0a79c",
 79}
 80
 81SPLIT_FOLDERS = {"cancer": "Cancer", "normal_mucosa": "Normal mucosa"}
 82
 83
 84class _RemoteZipReader:
 85    """Seekable file-like object over a remote zip, for reading its (small) structural data."""
 86
 87    def __init__(self, url, size):
 88        self.url = url
 89        self.size = size
 90        self.pos = 0
 91
 92    def seek(self, offset, whence=0):
 93        if whence == 0:
 94            self.pos = offset
 95        elif whence == 1:
 96            self.pos += offset
 97        elif whence == 2:
 98            self.pos = self.size + offset
 99        return self.pos
100
101    def tell(self):
102        return self.pos
103
104    def read(self, n=-1):
105        end = self.size - 1 if n is None or n < 0 else min(self.pos + n, self.size) - 1
106        if end < self.pos:
107            return b""
108        r = requests.get(self.url, headers={"Range": f"bytes={self.pos}-{end}"})
109        r.raise_for_status()
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    with requests.get(url, stream=True, allow_redirects=True) as r:
124        r.raise_for_status()
125        return int(r.headers["Content-Length"])
126
127
128def _list_zip_index(url):
129    """List a remote zip's central directory once, so individual members can be fetched by offset."""
130    size = _remote_size(url)
131    zf = zipfile.ZipFile(_RemoteZipReader(url, size))
132    return size, {info.filename: info for info in zf.infolist()}
133
134
135def _fetch_zip_member(url, size, info, dst_path):
136    if os.path.exists(dst_path):
137        return
138
139    # Fetch the local file header (with headroom for its variable-length fields) and the compressed
140    # payload in one range request. This dataset's members carry short names and no extra fields, so
141    # 256 bytes of headroom is always enough.
142    start = info.header_offset
143    end = start + 256 + info.compress_size - 1
144    r = requests.get(url, headers={"Range": f"bytes={start}-{end}"})
145    r.raise_for_status()
146    buf = r.content
147
148    fn_len, extra_len = struct.unpack("<HH", buf[26:30])
149    data_start = 30 + fn_len + extra_len
150    compressed = buf[data_start:data_start + info.compress_size]
151    if len(compressed) != info.compress_size:
152        raise RuntimeError(f"Incomplete read for zip member '{info.filename}'.")
153
154    data = compressed if info.compress_type == zipfile.ZIP_STORED else zlib.decompressobj(-15).decompress(compressed)
155
156    os.makedirs(os.path.dirname(dst_path), exist_ok=True)
157    tmp_path = dst_path + ".tmp"
158    with open(tmp_path, "wb") as f:
159        f.write(data)
160    os.replace(tmp_path, dst_path)
161
162
163def _resolve_stains(stains):
164    if stains is None:
165        return list(URLS)
166    if isinstance(stains, str):
167        stains = [stains]
168    invalid_stains = set(stains) - set(URLS)
169    if invalid_stains:
170        raise ValueError(f"Invalid stain choices: {sorted(invalid_stains)}. Choose from {sorted(URLS)}.")
171    return stains
172
173
174def _matches_split(member_name, split):
175    return split is None or f"/{SPLIT_FOLDERS[split]}/" in member_name
176
177
178def _matches_sample_ids(member_name, sample_ids):
179    if sample_ids is None:
180        return True
181    stem = Path(member_name).stem
182    return any(sample_id in stem for sample_id in sample_ids)
183
184
185def get_crc_epithelium_data(
186    path: Union[os.PathLike, str],
187    stains: Optional[Union[str, List[str]]] = None,
188    split: Optional[Literal["cancer", "normal_mucosa"]] = None,
189    sample_ids: Optional[List[str]] = None,
190    download: bool = False,
191) -> str:
192    """Download the CRC epithelium segmentation data.
193
194    Args:
195        path: Filepath to a folder where the downloaded data will be saved.
196        stains: The stain(s) to download. By default all 18 stains are downloaded.
197        split: The tissue split to restrict the data to. By default both splits are used.
198        sample_ids: The core ids to restrict the data to, e.g. ['A001-4']. By default all cores are used.
199        download: Whether to download the data if it is not present.
200
201    Returns:
202        Filepath to the folder where the raw images and masks are stored.
203    """
204    stains = _resolve_stains(stains)
205    path = str(path)
206    os.makedirs(path, exist_ok=True)
207
208    for stain in stains:
209        url = URLS[stain]
210        size, index = _list_zip_index(url)
211        raw_members = sorted(
212            name for name in index
213            if name.lower().endswith(".jpg") and _matches_split(name, split) and _matches_sample_ids(name, sample_ids)
214        )
215        if not raw_members:
216            raise RuntimeError(f"No members of stain '{stain}' match the requested 'split' / 'sample_ids'.")
217
218        for raw_name in tqdm(raw_members, desc=f"Fetch {stain} cores"):
219            mask_name = raw_name[:-len(".jpg")] + ".png"
220            if mask_name not in index:
221                raise RuntimeError(f"Missing mask '{mask_name}' for raw image '{raw_name}' in stain '{stain}'.")
222
223            raw_path = os.path.join(path, raw_name)
224            mask_path = os.path.join(path, mask_name)
225            if os.path.exists(raw_path) and os.path.exists(mask_path):
226                continue
227            if not download:
228                raise RuntimeError(f"Data for stain '{stain}' is not found and download is set to False.")
229
230            _fetch_zip_member(url, size, index[raw_name], raw_path)
231            _fetch_zip_member(url, size, index[mask_name], mask_path)
232
233    return path
234
235
236def get_crc_epithelium_paths(
237    path: Union[os.PathLike, str],
238    stains: Optional[Union[str, List[str]]] = None,
239    split: Optional[Literal["cancer", "normal_mucosa"]] = None,
240    sample_ids: Optional[List[str]] = None,
241    download: bool = False,
242) -> Tuple[List[str], List[str]]:
243    """Get paths to the CRC epithelium segmentation images and masks.
244
245    Args:
246        path: Filepath to a folder where the downloaded data will be saved.
247        stains: The stain(s) to use. By default all 18 stains are used.
248        split: The tissue split to restrict the data to. By default both splits are used.
249        sample_ids: The core ids to restrict the data to, e.g. ['A001-4']. By default all cores are used.
250        download: Whether to download the data if it is not present.
251
252    Returns:
253        List of filepaths for the image data.
254        List of filepaths for the label data.
255    """
256    stains = _resolve_stains(stains)
257    data_dir = get_crc_epithelium_data(path, stains, split, sample_ids, download)
258
259    raw_paths, label_paths = [], []
260    for stain in stains:
261        stain_raw_paths = sorted(glob(os.path.join(data_dir, stain, "**", "*.jpg"), recursive=True))
262        for raw_path in stain_raw_paths:
263            member_name = os.path.relpath(raw_path, data_dir)
264            if not (_matches_split(member_name, split) and _matches_sample_ids(member_name, sample_ids)):
265                continue
266            mask_path = os.path.splitext(raw_path)[0] + ".png"
267            if not os.path.exists(mask_path):
268                raise RuntimeError(f"Missing mask for raw image '{raw_path}'.")
269            raw_paths.append(raw_path)
270            label_paths.append(mask_path)
271
272    if not raw_paths:
273        raise RuntimeError("Could not find any images and masks for the requested settings.")
274
275    return raw_paths, label_paths
276
277
278def get_crc_epithelium_dataset(
279    path: Union[os.PathLike, str],
280    patch_shape: Tuple[int, int],
281    stains: Optional[Union[str, List[str]]] = None,
282    split: Optional[Literal["cancer", "normal_mucosa"]] = None,
283    sample_ids: Optional[List[str]] = None,
284    resize_inputs: bool = False,
285    download: bool = False,
286    **kwargs,
287) -> Dataset:
288    """Get the CRC epithelium segmentation dataset.
289
290    Args:
291        path: Filepath to a folder where the downloaded data will be saved.
292        patch_shape: The patch shape to use for training.
293        stains: The stain(s) to use. By default all 18 stains are used.
294        split: The tissue split to restrict the data to. By default both splits are used.
295        sample_ids: The core ids to restrict the data to, e.g. ['A001-4']. By default all cores are used.
296        resize_inputs: Whether to resize the inputs.
297        download: Whether to download the data if it is not present.
298        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
299
300    Returns:
301        The segmentation dataset.
302    """
303    raw_paths, label_paths = get_crc_epithelium_paths(path, stains, split, sample_ids, download)
304
305    if resize_inputs:
306        resize_kwargs = {"patch_shape": patch_shape, "is_rgb": True}
307        kwargs, patch_shape = util.update_kwargs_for_resize_trafo(
308            kwargs=kwargs, patch_shape=patch_shape, resize_inputs=resize_inputs, resize_kwargs=resize_kwargs
309        )
310
311    return torch_em.default_segmentation_dataset(
312        raw_paths=raw_paths,
313        raw_key=None,
314        label_paths=label_paths,
315        label_key=None,
316        patch_shape=patch_shape,
317        is_seg_dataset=False,
318        ndim=2,
319        with_channels=True,
320        **kwargs,
321    )
322
323
324def get_crc_epithelium_loader(
325    path: Union[os.PathLike, str],
326    batch_size: int,
327    patch_shape: Tuple[int, int],
328    stains: Optional[Union[str, List[str]]] = None,
329    split: Optional[Literal["cancer", "normal_mucosa"]] = None,
330    sample_ids: Optional[List[str]] = None,
331    resize_inputs: bool = False,
332    download: bool = False,
333    **kwargs,
334) -> DataLoader:
335    """Get the CRC epithelium segmentation dataloader.
336
337    Args:
338        path: Filepath to a folder where the downloaded data will be saved.
339        batch_size: The batch size for training.
340        patch_shape: The patch shape to use for training.
341        stains: The stain(s) to use. By default all 18 stains are used.
342        split: The tissue split to restrict the data to. By default both splits are used.
343        sample_ids: The core ids to restrict the data to, e.g. ['A001-4']. By default all cores are used.
344        resize_inputs: Whether to resize the inputs.
345        download: Whether to download the data if it is not present.
346        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or the PyTorch DataLoader.
347
348    Returns:
349        The DataLoader.
350    """
351    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
352    dataset = get_crc_epithelium_dataset(
353        path, patch_shape, stains, split, sample_ids, resize_inputs, download, **ds_kwargs
354    )
355    return torch_em.get_data_loader(dataset, batch_size, **loader_kwargs)