torch_em.data.datasets.histopathology.spatch

  1"""SPATCH contains expert annotations for nucleus segmentation in images of human tumor tissue.
  2
  3The dataset covers three tumor types (ovarian cancer, hepatocellular carcinoma and colon adenocarcinoma)
  4imaged on four spatial transcriptomics platforms. Xenium and CosMx provide DAPI images, while Visium HD
  5and Stereo-seq provide H&E images. Every subset ships five manually annotated tiles.
  6NOTE: The annotations mark nuclear boundaries. They come as polygons, which this module rasterizes.
  7
  8The dataset is hosted at https://spatch.pku-genomics.org.
  9This dataset is from the publication https://doi.org/10.1038/s41467-025-64292-3.
 10Please cite it if you use this dataset for your research.
 11"""
 12
 13import os
 14import json
 15import zipfile
 16import urllib.request
 17from glob import glob
 18from natsort import natsorted
 19from typing import Union, Tuple, List
 20
 21import numpy as np
 22import imageio.v3 as imageio
 23from skimage.draw import polygon
 24from skimage.segmentation import relabel_sequential
 25
 26from torch.utils.data import Dataset, DataLoader
 27
 28import torch_em
 29
 30from .. import util
 31
 32
 33TOKEN_URL = "https://bj21400.api.aliyunfile.com/v2/share_link/get_share_token"
 34LIST_URL = "https://bj21400.api.aliyunfile.com/v2/file/list"
 35
 36# The portal serves every archive through a share link, which is addressed by the id below.
 37SHARE_IDS = {
 38    "xenium_ov": "PXSYun7LWov",
 39    "xenium_hcc": "mzaFwMVEuvq",
 40    "xenium_coad": "TUzo1AaZ9Vx",
 41    "cosmx_ov": "ttidbCgTY1z",
 42    "cosmx_hcc": "aeubu9Uxt2Z",
 43    "cosmx_coad": "Ja4mTzWciZb",
 44    "visium_hd_ov": "Joubkps1Mz1",
 45    "visium_hd_hcc": "sY36AC8GoPW",
 46    "visium_hd_coad": "PjBvUPtFaqP",
 47    "stereoseq_ov": "3Es7YQz69ja",
 48}
 49
 50CHECKSUMS = {
 51    "xenium_ov": "674a7e317c75d5718bd42d8fdcc11fd1d0f9f7274b597e074c453a14de83452e",
 52    "xenium_hcc": "256052bfc589be0c765d69ec884568dad54b1e395e32943a155388a087f93f72",
 53    "xenium_coad": "a3c04a7f8d307508e8f4382483341e6fb69f3f4c4630ca54f43c6ee64792c4b5",
 54    "cosmx_ov": "040fd2d29a0fc278c0ef46c231d6825c049c6c87feb31c5a6d96bdaf9ff09d37",
 55    "cosmx_hcc": "39dccac329d50218524429002c246681310dfb15091eaa8431fafe8d1c318aa3",
 56    "cosmx_coad": "dda1af15d0e2c191a83d9f518661a420cb51716e428b860612a60d5155e5ec2f",
 57    "visium_hd_ov": "3e25df5fb9964a95158810bbb621b988fdcaa81de428d9a5f6c34ee0aa6099e7",
 58    "visium_hd_hcc": "322eda304ec22e49037b35bd3d60c3d565c4dd202609c5c570a05f4672beedb9",
 59    "visium_hd_coad": "7aa9e5da53fc8aa83199f10383e945aa85f147e4be3d8db61b824d13cbe05441",
 60    "stereoseq_ov": "4187167c8d622f8f187d5d0f30bb3ce599154e2b01d8e68147faa9042b12e88d",
 61}
 62
 63
 64def _post_json(url, payload, headers=None):
 65    request = urllib.request.Request(url, data=json.dumps(payload).encode(), method="POST")
 66    request.add_header("Content-Type", "application/json")
 67    for key, value in (headers or {}).items():
 68        request.add_header(key, value)
 69    with urllib.request.urlopen(request, timeout=60) as response:
 70        return json.loads(response.read())
 71
 72
 73def _get_download_url(share_id):
 74    """Resolve the share id of a subset to a download url, which the portal only issues on request."""
 75    token = _post_json(TOKEN_URL, {"share_id": share_id, "ignoreError": True})["share_token"]
 76    payload = {
 77        "limit": 100,
 78        "marker": "",
 79        "share_id": share_id,
 80        "parent_file_id": "root",
 81        "fields": "user_name,dir_size,url,content_type,upload_id,crc64_hash,revision_id,description",
 82        "url_expire_sec": 7200,
 83    }
 84    items = _post_json(LIST_URL, payload, {"x-share-token": token})["items"]
 85    if not items:
 86        raise RuntimeError(f"The share link '{share_id}' does not contain any file.")
 87    return items[0]["download_url"]
 88
 89
 90def _get_outlines(json_path):
 91    """Return the nucleus outlines, which come in the labelme format, or in the darwin format.
 92
 93    Some tiles annotate the same nucleus twice, so this drops the repeated outlines.
 94    """
 95    with open(json_path) as f:
 96        content = json.load(f)
 97
 98    if "annotations" in content:  # darwin
 99        paths = [path for a in content["annotations"] for path in a.get("polygon", {}).get("paths", [])]
100        outlines = [np.array([[p["y"], p["x"]] for p in path]) for path in paths]
101    else:
102        shapes = [s for s in content["shapes"] if s.get("shape_type") == "polygon"]
103        outlines = [np.array([[p[1], p[0]] for p in s["points"]]) for s in shapes]
104
105    seen, unique = set(), []
106    for outline in outlines:
107        key = np.round(outline, 4).tobytes()
108        if key not in seen:
109            seen.add(key)
110            unique.append(outline)
111    return unique
112
113
114def _get_instances(json_path, shape):
115    """Paint the outlines into a label image.
116
117    Nuclei that the annotators drew on top of each other cannot all survive this, so the labels
118    are made consecutive afterwards.
119    """
120    outlines = _get_outlines(json_path)
121    instances = np.zeros(shape, dtype="uint16")
122    for label, outline in enumerate(outlines, start=1):
123        rr, cc = polygon(outline[:, 0], outline[:, 1], shape=shape)
124        instances[rr, cc] = label
125    return relabel_sequential(instances)[0]
126
127
128def _preprocess_data(input_dir, data_dir):
129    import h5py
130
131    os.makedirs(data_dir, exist_ok=True)
132    tile_paths = natsorted(glob(os.path.join(input_dir, "**", "tile*.png"), recursive=True))
133    if not tile_paths:
134        raise RuntimeError(f"Could not find the annotated tiles in {input_dir}.")
135
136    for tile_path in tile_paths:
137        tile = os.path.splitext(os.path.basename(tile_path))[0]
138        json_path = os.path.join(os.path.dirname(tile_path), f"{tile.replace('tile', 'mask')}.json")
139        out_path = os.path.join(data_dir, f"{tile}.h5")
140        if os.path.exists(out_path) or not os.path.exists(json_path):
141            continue
142
143        image = imageio.imread(tile_path)[..., :3]
144        instances = _get_instances(json_path, image.shape[:2])
145
146        with h5py.File(out_path, "a") as f:
147            f.create_dataset("raw/rgb", data=image.transpose(2, 0, 1), compression="gzip")
148            f.create_dataset("labels/nuclei", data=instances, compression="gzip")
149
150
151def get_spatch_data(path: Union[os.PathLike, str], subset: str, download: bool = False) -> str:
152    """Download one subset of the SPATCH dataset.
153
154    Args:
155        path: The folder where the function stores the data.
156        subset: The subset of the dataset. See `SHARE_IDS` for the valid choices.
157        download: Whether to download the data if it is not present.
158
159    Returns:
160        The filepath to the folder with the prepared data.
161    """
162    if subset not in SHARE_IDS:
163        raise ValueError(f"'{subset}' is not a valid subset. Choose one of {list(SHARE_IDS.keys())}.")
164
165    data_dir = os.path.join(path, subset, "data")
166    if glob(os.path.join(data_dir, "*.h5")):
167        return data_dir
168
169    subset_dir = os.path.join(path, subset)
170    os.makedirs(subset_dir, exist_ok=True)
171
172    input_dir = os.path.join(subset_dir, "manual_segmentation")
173    if not os.path.exists(input_dir):
174        zip_path = os.path.join(subset_dir, f"{subset}.zip")
175        if not os.path.exists(zip_path):
176            if not download:
177                raise RuntimeError(f"The data for '{subset}' is not present and download is set to False.")
178            util.download_source(
179                path=zip_path, url=_get_download_url(SHARE_IDS[subset]), download=True, checksum=CHECKSUMS[subset]
180            )
181        with zipfile.ZipFile(zip_path) as f:
182            f.extractall(subset_dir)
183
184    _preprocess_data(input_dir, data_dir)
185
186    return data_dir
187
188
189def get_spatch_paths(
190    path: Union[os.PathLike, str], subset: Union[str, List[str]], download: bool = False
191) -> List[str]:
192    """Get the paths to the SPATCH data.
193
194    Args:
195        path: The folder where the function stores the data.
196        subset: One subset or a list of subsets. See `SHARE_IDS` for the valid choices.
197        download: Whether to download the data if it is not present.
198
199    Returns:
200        The list of filepaths to the input data.
201    """
202    subsets = [subset] if isinstance(subset, str) else subset
203    volume_paths = []
204    for name in subsets:
205        data_dir = get_spatch_data(path, name, download)
206        volume_paths.extend(natsorted(glob(os.path.join(data_dir, "*.h5"))))
207
208    assert len(volume_paths) > 0, f"Could not find data for the subset '{subset}'."
209    return volume_paths
210
211
212def get_spatch_dataset(
213    path: Union[os.PathLike, str],
214    patch_shape: Tuple[int, int],
215    subset: Union[str, List[str]],
216    download: bool = False,
217    **kwargs
218) -> Dataset:
219    """Get the SPATCH dataset for nucleus segmentation in images of human tumor tissue.
220
221    Args:
222        path: The folder where the function stores the data.
223        patch_shape: The patch shape to use for training.
224        subset: One subset or a list of subsets. See `SHARE_IDS` for the valid choices.
225        download: Whether to download the data if it is not present.
226        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
227
228    Returns:
229        The segmentation dataset.
230    """
231    volume_paths = get_spatch_paths(path, subset, download)
232    kwargs = util.update_kwargs(kwargs, "with_channels", True)
233
234    return torch_em.default_segmentation_dataset(
235        raw_paths=volume_paths,
236        raw_key="raw/rgb",
237        label_paths=volume_paths,
238        label_key="labels/nuclei",
239        patch_shape=patch_shape,
240        is_seg_dataset=True,
241        ndim=2,
242        **kwargs
243    )
244
245
246def get_spatch_loader(
247    path: Union[os.PathLike, str],
248    batch_size: int,
249    patch_shape: Tuple[int, int],
250    subset: Union[str, List[str]],
251    download: bool = False,
252    **kwargs
253) -> DataLoader:
254    """Get the SPATCH dataloader for nucleus segmentation in images of human tumor tissue.
255
256    Args:
257        path: The folder where the function stores the data.
258        batch_size: The batch size for training.
259        patch_shape: The patch shape to use for training.
260        subset: One subset or a list of subsets. See `SHARE_IDS` for the valid choices.
261        download: Whether to download the data if it is not present.
262        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or for the PyTorch DataLoader.
263
264    Returns:
265        The DataLoader.
266    """
267    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
268    dataset = get_spatch_dataset(path, patch_shape, subset, download, **ds_kwargs)
269    return torch_em.get_data_loader(dataset, batch_size=batch_size, **loader_kwargs)