torch_em.data.datasets.light_microscopy.xenium

  1"""The Xenium datasets contain cell and nucleus segmentation masks for whole-slide
  2fluorescence images of human and mouse tissue, from the 10x Genomics in situ platform.
  3
  4Every sample pairs a set of morphology images with the masks that the Xenium Onboard
  5Analysis (XOA) pipeline produced. Channel 0 always holds the DAPI nuclear stain. The
  6other channels hold the boundary and interior stains of the multi-tissue stain mix,
  7which XOA uses to draw the cell outlines.
  8
  9NOTE: The masks are the output of XOA, not a manual annotation. XOA segments the nuclei
 10on DAPI, then grows each cell from its nucleus using the boundary and interior stains.
 11Where no stain supports a boundary it falls back to expanding the nucleus by 5 micrometer,
 12so a share of the cells are dilated nuclei rather than observed outlines. It runs between three
 13and five percent on the samples here. XOA reports that share from version 3.0 on, and each h5
 14file passes it through as the 'nuc_expansion_fraction' attribute, which stays 'nan' for the
 15earlier versions. Treat this data as a weak label source for pretraining, not as a benchmark.
 16For expert annotations on Xenium images see the SPATCH dataset in
 17`torch_em.data.datasets.histopathology.spatch`.
 18
 19NOTE: A cell can hold more than one nucleus, so the two label sets do not match one to one
 20and their ids do not correspond.
 21
 22NOTE: The layout of the bundle changed with the XOA version. Version 2 and 3 name the
 23morphology channels 'morphology_focus_000N.ome.tif' and version 4 names them
 24'ch000N_<stain>.ome.tif'. This loader sorts the files by name, which puts DAPI first in
 25both schemes, and stores the names it used in the 'channel_files' attribute. Versions
 26before 2.0 ship a different layout and this loader does not cover them.
 27
 28NOTE: The images are whole slides that hold large empty regions. Pass a sampler such as
 29`torch_em.data.MinInstanceSampler` to avoid drawing empty patches.
 30
 31The data is hosted at https://www.10xgenomics.com/datasets under the CC BY 4.0 license.
 32Please cite the corresponding dataset page if you use this data in your research.
 33"""
 34
 35import os
 36import json
 37import zlib
 38import shutil
 39import struct
 40import hashlib
 41import zipfile
 42from glob import glob
 43from natsort import natsorted
 44from contextlib import contextmanager
 45from typing import Dict, List, Literal, Optional, Sequence, Tuple, Union
 46
 47import requests
 48
 49from torch.utils.data import DataLoader, Dataset
 50
 51import torch_em
 52
 53from .. import util
 54
 55
 56URLS = {
 57    "human_pancreas": "https://cf.10xgenomics.com/samples/xenium/2.0.0/Xenium_V1_human_Pancreas_FFPE/Xenium_V1_human_Pancreas_FFPE_xe_outs.zip",  # noqa
 58    "human_lung_cancer": "https://cf.10xgenomics.com/samples/xenium/2.0.0/Xenium_V1_humanLung_Cancer_FFPE/Xenium_V1_humanLung_Cancer_FFPE_xe_outs.zip",  # noqa
 59    "mouse_colon": "https://cf.10xgenomics.com/samples/xenium/2.0.0/Xenium_V1_mouse_Colon_FF/Xenium_V1_mouse_Colon_FF_xe_outs.zip",  # noqa
 60    "human_skin": "https://cf.10xgenomics.com/samples/xenium/3.0.0/Xenium_Prime_Human_Skin_FFPE/Xenium_Prime_Human_Skin_FFPE_xe_outs.zip",  # noqa
 61    "human_prostate": "https://cf.10xgenomics.com/samples/xenium/3.0.0/Xenium_Prime_Human_Prostate_FFPE/Xenium_Prime_Human_Prostate_FFPE_xe_outs.zip",  # noqa
 62    "human_breast": "https://cf.10xgenomics.com/samples/xenium/4.0.0/Human_Breast_Biomarkers_S1_Top/Human_Breast_Biomarkers_S1_Top_xe_outs.zip",  # noqa
 63}
 64
 65# The sha256 of the two members that every download reads, keyed by the member name.
 66# The optional stain channels carry no entry here and rest on their crc32 alone.
 67CHECKSUMS = {
 68    "human_pancreas": {
 69        "cells.zarr.zip": "5747730a532f130b41f83e970784302623a784dfc9fe571dbd11e9179366387c",
 70        "morphology_focus/morphology_focus_0000.ome.tif":
 71            "3a7c669ebe6cbbd7e90ebab790bc8e666b9165e972bc528e134037eef632b3bd",
 72    },
 73    "human_lung_cancer": {
 74        "cells.zarr.zip": "f274590902f9162295d292e540b631ba37c03574a8db754a5c27f101c8c94ef1",
 75        "morphology_focus/morphology_focus_0000.ome.tif":
 76            "eb0ef6aa2c7adfa5e70323b4b9530771f8679bd9c0eb6dbd1af40faabc2ca398",
 77    },
 78    "mouse_colon": {
 79        "cells.zarr.zip": "cc5f811511d349955b15b6e4ebb4205a4208418f900ada5853231918d929a17e",
 80        "morphology_focus/morphology_focus_0000.ome.tif":
 81            "2a723a59b01c943d98a31aa47fc0297e6e9414b391c1d4d3949df0b27e510ef2",
 82    },
 83    "human_skin": {
 84        "cells.zarr.zip": "ef805ec5c7e733f7c82a0e88652a04550abfd9f7a6c60c52339aa34ff780316d",
 85        "morphology_focus/morphology_focus_0000.ome.tif":
 86            "435caf0d766d6226ab8a8c3be551c9a87ad2be1997d37767c992aa36bd7d09f3",
 87    },
 88    "human_prostate": {
 89        "cells.zarr.zip": "7ce2e92b1b085b35051f8568a983c8f8a8622e73bfae6060efdf7b47de6371e4",
 90        "morphology_focus/morphology_focus_0000.ome.tif":
 91            "a095080cc946f734545db51300d7ba24700c03b02cba48204e680c556dd867ed",
 92    },
 93    "human_breast": {
 94        "cells.zarr.zip": "61d09dd84509a09cf2f6c1252e7d12772f6b3d6f33dd02ea25555f8ad84b6758",
 95        "morphology_focus/ch0000_dapi.ome.tif":
 96            "75476f9be2b0f936022611d79717cb77356e90ec40902531098b984a54d9c4a0",
 97    },
 98}
 99
100# The tissue of every sample, for the h5 attributes.
101TISSUES = {
102    "human_pancreas": "human pancreas",
103    "human_lung_cancer": "human lung adenocarcinoma",
104    "mouse_colon": "mouse colon",
105    "human_skin": "human skin melanoma",
106    "human_prostate": "human prostate adenocarcinoma",
107    "human_breast": "human breast",
108}
109
110LABEL_CHANNELS = {"nuclei": "0", "cells": "1"}
111
112BLOCK = 4096
113
114
115def _read_range(url: str, start: int, end: int) -> bytes:
116    """Read the bytes from `start` to `end`, both included."""
117    response = requests.get(url, headers={"Range": f"bytes={start}-{end}"}, stream=True)
118    response.raise_for_status()
119    return response.content
120
121
122def _read_zip64_extra(extra: bytes, uncompressed: int, compressed: int, offset: int) -> Tuple[int, int, int]:
123    """Replace the fields that overflowed the 32 bit form with their zip64 values."""
124    position = 0
125    while position + 4 <= len(extra):
126        tag, size = struct.unpack("<HH", extra[position:position + 4])
127        if tag == 0x0001:
128            values = list(struct.unpack(f"<{size // 8}Q", extra[position + 4:position + 4 + size]))
129            if uncompressed == 0xFFFFFFFF:
130                uncompressed = values.pop(0)
131            if compressed == 0xFFFFFFFF:
132                compressed = values.pop(0)
133            if offset == 0xFFFFFFFF:
134                offset = values.pop(0)
135            break
136        position += 4 + size
137    return uncompressed, compressed, offset
138
139
140def _remote_zip_index(url: str) -> Dict[str, Dict[str, int]]:
141    """Read the central directory of a remote zip, which lists every member and where it sits.
142
143    This turns the bundle into a set of files that can be read one by one, so that a download
144    can skip the transcripts and the expression matrix, which segmentation never touches.
145    """
146    response = requests.head(url)
147    response.raise_for_status()
148    size = int(response.headers["Content-Length"])
149
150    tail = _read_range(url, max(0, size - 262144), size - 1)
151    position = tail.rfind(b"PK\x06\x06")
152    if position == -1:
153        position = tail.rfind(b"PK\x05\x06")
154        if position == -1:
155            raise RuntimeError(f"Could not find the central directory of the zip at {url}.")
156        directory_size, directory_offset = struct.unpack("<II", tail[position + 12:position + 20])
157    else:
158        directory_size, directory_offset = struct.unpack("<QQ", tail[position + 40:position + 56])
159
160    directory = _read_range(url, directory_offset, directory_offset + directory_size + 128)
161    index, position = {}, 0
162    while position + 46 <= len(directory) and directory[position:position + 4] == b"PK\x01\x02":
163        method = struct.unpack("<H", directory[position + 10:position + 12])[0]
164        crc = struct.unpack("<I", directory[position + 16:position + 20])[0]
165        compressed, uncompressed = struct.unpack("<II", directory[position + 20:position + 28])
166        name_length, extra_length, comment_length = struct.unpack("<HHH", directory[position + 28:position + 34])
167        offset = struct.unpack("<I", directory[position + 42:position + 46])[0]
168        name = directory[position + 46:position + 46 + name_length].decode("utf-8", errors="replace")
169        extra = directory[position + 46 + name_length:position + 46 + name_length + extra_length]
170        uncompressed, compressed, offset = _read_zip64_extra(extra, uncompressed, compressed, offset)
171        index[name] = {"method": method, "crc": crc, "compressed": compressed, "offset": offset}
172        position += 46 + name_length + extra_length + comment_length
173    return index
174
175
176def _download_member(url: str, name: str, entry: Dict[str, int], output_path: str, expected_sha: Optional[str]) -> None:
177    """Read one member out of a remote zip and write it, checking its crc32 and its sha256."""
178    from tqdm import tqdm
179
180    # The local header repeats the name and carries its own extra field, so the data of the
181    # member starts behind a header whose length only the header itself gives.
182    header = _read_range(url, entry["offset"], entry["offset"] + 29)
183    if header[:4] != b"PK\x03\x04":
184        raise RuntimeError(f"The member '{name}' does not start with a local header.")
185    name_length, extra_length = struct.unpack("<HH", header[26:30])
186    start = entry["offset"] + 30 + name_length + extra_length
187
188    if entry["method"] not in (0, 8):
189        raise RuntimeError(f"The member '{name}' uses the unsupported compression method {entry['method']}.")
190    decompressor = zlib.decompressobj(-zlib.MAX_WBITS) if entry["method"] == 8 else None
191
192    crc, digest = 0, hashlib.sha256()
193    temporary_path = f"{output_path}.incomplete"
194    os.makedirs(os.path.dirname(output_path), exist_ok=True)
195
196    headers = {"Range": f"bytes={start}-{start + entry['compressed'] - 1}"}
197    with requests.get(url, headers=headers, stream=True) as response:
198        response.raise_for_status()
199        description = f"Download {name}"
200        with open(temporary_path, "wb") as f:
201            with tqdm(total=entry["compressed"], unit="B", unit_scale=True, desc=description) as progress:
202                for chunk in response.iter_content(chunk_size=1 << 20):
203                    progress.update(len(chunk))
204                    block = decompressor.decompress(chunk) if decompressor is not None else chunk
205                    crc = zlib.crc32(block, crc)
206                    digest.update(block)
207                    f.write(block)
208            if decompressor is not None:
209                block = decompressor.flush()
210                crc = zlib.crc32(block, crc)
211                digest.update(block)
212                f.write(block)
213
214    if crc != entry["crc"]:
215        os.remove(temporary_path)
216        raise RuntimeError(f"The member '{name}' has the crc32 {crc}, but the zip lists {entry['crc']}.")
217    if expected_sha is not None and digest.hexdigest() != expected_sha:
218        os.remove(temporary_path)
219        raise RuntimeError(f"The member '{name}' has the sha256 {digest.hexdigest()}, expected {expected_sha}.")
220    os.replace(temporary_path, output_path)
221
222
223def _download_bundle(url: str, bundle_dir: str, sample: str, download: bool, with_stains: bool) -> str:
224    """Download the members that this loader reads, and skip the rest of the bundle."""
225    if os.path.exists(bundle_dir):
226        return bundle_dir
227    if not download:
228        raise RuntimeError(f"Cannot find the data at {bundle_dir}, but download was set to False")
229
230    index = _remote_zip_index(url)
231    channels = natsorted(n for n in index if n.startswith("morphology_focus/") and n.endswith(".ome.tif"))
232    if not channels:
233        raise RuntimeError(f"The bundle for '{sample}' holds no morphology channel.")
234    if not with_stains:
235        channels = channels[:1]
236
237    names = ["experiment.xenium", "cells.zarr.zip"] + channels
238    missing = [n for n in names if n not in index]
239    if missing:
240        raise RuntimeError(f"The bundle for '{sample}' is missing {missing}.")
241
242    temporary_dir = f"{bundle_dir}.tmp"
243    if os.path.exists(temporary_dir):
244        shutil.rmtree(temporary_dir)
245    for name in names:
246        _download_member(url, name, index[name], os.path.join(temporary_dir, name), CHECKSUMS[sample].get(name))
247
248    os.replace(temporary_dir, bundle_dir)
249    return bundle_dir
250
251
252def _channel_paths(bundle_dir: str) -> List[str]:
253    """Return the morphology channels, DAPI first. Both naming schemes sort DAPI to the front."""
254    paths = natsorted(glob(os.path.join(bundle_dir, "morphology_focus", "*.ome.tif")))
255    if not paths:
256        raise RuntimeError(f"Could not find any morphology channel in {bundle_dir}.")
257    return paths
258
259
260@contextmanager
261def _open_image(path: str):
262    """Open one morphology channel at full resolution, without reading it into memory.
263
264    The channels form a multi file OME pyramid, which the OME reader of tifffile rejects.
265    Reading the file on its own gives the pyramid of this channel alone.
266    """
267    import zarr
268    import tifffile
269
270    with tifffile.TiffFile(path, is_ome=False) as tif:
271        store = tif.series[0].levels[0].aszarr()
272        try:
273            yield zarr.open(store, mode="r")["0"]
274        finally:
275            store.close()
276
277
278def _copy_blockwise(source, target, channel: Optional[int] = None) -> int:
279    """Copy a large 2d array block by block, to keep the memory use bounded.
280
281    The target must be indexed in one step, because indexing an h5 dataset twice writes
282    into a copy of the data rather than into the file.
283
284    Returns the largest value that was copied, which gives the instance count of a mask
285    without a second pass over the whole array.
286    """
287    height, width = source.shape[-2:]
288    largest = 0
289    for y in range(0, height, BLOCK):
290        for x in range(0, width, BLOCK):
291            box = (slice(y, min(y + BLOCK, height)), slice(x, min(x + BLOCK, width)))
292            block = source[box]
293            if channel is None:
294                target[box] = block
295            else:
296                target[(channel,) + box] = block
297            largest = max(largest, int(block.max()))
298    return largest
299
300
301def _create_h5(bundle_dir: str, output_path: str, sample: str, with_stains: bool) -> str:
302    import h5py
303    import zarr
304    from tqdm import tqdm
305
306    if os.path.exists(output_path):
307        return output_path
308
309    with open(os.path.join(bundle_dir, "experiment.xenium")) as f:
310        manifest = json.load(f)
311
312    masks_zip = os.path.join(bundle_dir, "cells.zarr.zip")
313    masks_dir = os.path.join(bundle_dir, "cells.zarr")
314    if not os.path.exists(masks_dir):
315        with zipfile.ZipFile(masks_zip) as archive:
316            archive.extractall(masks_dir)
317    masks = zarr.open(masks_dir, mode="r")["masks"]
318
319    channel_paths = _channel_paths(bundle_dir)
320    with _open_image(channel_paths[0]) as dapi:
321        shape = dapi.shape
322    for name, key in LABEL_CHANNELS.items():
323        if masks[key].shape != shape:
324            raise RuntimeError(
325                f"The {name} mask of '{sample}' has the shape {masks[key].shape}, "
326                f"but its DAPI image has the shape {shape}."
327            )
328
329    temporary_path = f"{output_path}.tmp"
330    with h5py.File(temporary_path, "w") as f:
331        f.attrs["sample"] = sample
332        f.attrs["tissue"] = TISSUES[sample]
333        f.attrs["axes"] = "yx"
334        f.attrs["pixel_size"] = manifest["pixel_size"]
335        f.attrs["xoa_version"] = manifest["analysis_sw_version"]
336        f.attrs["segmentation_stain"] = manifest.get("segmentation_stain", "none")
337        f.attrs["nuc_expansion_fraction"] = manifest.get("segmented_cell_nuc_expansion_frac", float("nan"))
338        f.attrs["num_cells"] = manifest["num_cells"]
339        f.attrs["channel_files"] = [os.path.basename(p) for p in channel_paths]
340
341        chunks = (min(1024, shape[0]), min(1024, shape[1]))
342        raw = f.create_dataset("raw/dapi", shape=shape, dtype="uint16", chunks=chunks, compression="gzip")
343        raw.attrs["pixel_size"] = manifest["pixel_size"]
344        with _open_image(channel_paths[0]) as dapi:
345            _copy_blockwise(dapi, raw)
346
347        if with_stains:
348            stack_shape = (len(channel_paths),) + shape
349            stack = f.create_dataset(
350                "raw/stack", shape=stack_shape, dtype="uint16", chunks=(1,) + chunks, compression="gzip"
351            )
352            stack.attrs["pixel_size"] = manifest["pixel_size"]
353            for channel, channel_path in enumerate(tqdm(channel_paths, desc=f"Copy channels of '{sample}'")):
354                with _open_image(channel_path) as image:
355                    if image.shape != shape:
356                        raise RuntimeError(
357                            f"The channel {os.path.basename(channel_path)} of '{sample}' has the shape "
358                            f"{image.shape}, but its DAPI image has the shape {shape}."
359                        )
360                    _copy_blockwise(image, stack, channel=channel)
361
362        for name, key in LABEL_CHANNELS.items():
363            labels = f.create_dataset(
364                f"labels/{name}", shape=shape, dtype="uint32", chunks=chunks, compression="gzip"
365            )
366            labels.attrs["pixel_size"] = manifest["pixel_size"]
367            labels.attrs["num_instances"] = _copy_blockwise(masks[key], labels)
368
369    os.replace(temporary_path, output_path)
370    shutil.rmtree(masks_dir, ignore_errors=True)
371    return output_path
372
373
374def get_xenium_data(
375    path: Union[os.PathLike, str],
376    sample: str,
377    download: bool = False,
378    with_stains: bool = True,
379) -> str:
380    """Download one Xenium sample.
381
382    The source is the Xenium Explorer subset of the output bundle, which holds the masks and
383    the morphology images but not the transcripts. This function does not fetch the whole
384    archive. It reads the central directory of the remote zip, then reads only the members
385    that it needs, which cuts the download by about four times without the stains.
386
387    Args:
388        path: Filepath to a folder where the downloaded data will be saved.
389        sample: The sample to use. See `URLS` for the available samples.
390        download: Whether to download the data if it is not present.
391        with_stains: Whether to also store the boundary and interior stains next to DAPI.
392            The cell masks rest on those stains, so a model that predicts cells needs them.
393
394    Returns:
395        The filepath to the h5 file of the sample.
396    """
397    if sample not in URLS:
398        raise ValueError(f"'{sample}' is not a valid sample. Choose from {list(URLS)}.")
399
400    output_path = os.path.join(path, "preprocessed", f"{sample}.h5")
401    if os.path.exists(output_path):
402        return output_path
403
404    os.makedirs(os.path.join(path, "preprocessed"), exist_ok=True)
405    bundle_dir = _download_bundle(URLS[sample], os.path.join(path, "bundles", sample), sample, download, with_stains)
406    _create_h5(bundle_dir, output_path, sample, with_stains)
407    shutil.rmtree(bundle_dir, ignore_errors=True)
408
409    return output_path
410
411
412def get_xenium_paths(
413    path: Union[os.PathLike, str],
414    sample: Optional[Union[str, Sequence[str]]] = None,
415    download: bool = False,
416    with_stains: bool = True,
417) -> List[str]:
418    """Get paths to the Xenium data.
419
420    Args:
421        path: Filepath to a folder where the downloaded data will be saved.
422        sample: The sample or samples to use. Defaults to all of them.
423        download: Whether to download the data if it is not present.
424        with_stains: Whether to also store the boundary and interior stains next to DAPI.
425
426    Returns:
427        List of filepaths for the h5 data.
428    """
429    if sample is None:
430        samples = list(URLS)
431    else:
432        samples = [sample] if isinstance(sample, str) else list(sample)
433
434    return [get_xenium_data(path, name, download, with_stains) for name in samples]
435
436
437def get_xenium_dataset(
438    path: Union[os.PathLike, str],
439    patch_shape: Tuple[int, int],
440    sample: Optional[Union[str, Sequence[str]]] = None,
441    label_channel: Literal["nuclei", "cells"] = "nuclei",
442    raw_channel: Literal["dapi", "stack"] = "dapi",
443    offsets: Optional[List[List[int]]] = None,
444    boundaries: bool = False,
445    binary: bool = False,
446    download: bool = False,
447    **kwargs,
448) -> Dataset:
449    """Get the Xenium dataset for cell and nucleus segmentation.
450
451    Args:
452        path: Filepath to a folder where the downloaded data will be saved.
453        patch_shape: The 2D patch shape to use for training.
454        sample: The sample or samples to use. Defaults to all of them.
455        label_channel: The masks to use as target. Either 'nuclei' or 'cells'.
456        raw_channel: The images to use as input. Either 'dapi' for the nuclear stain alone,
457            or 'stack' for all morphology channels.
458        offsets: Offset values for affinity computation used as target.
459        boundaries: Whether to compute boundaries as the target.
460        binary: Whether to use a binary segmentation target.
461        download: Whether to download the data if it is not present.
462        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset`.
463
464    Returns:
465        The segmentation dataset.
466    """
467    if len(patch_shape) != 2:
468        raise ValueError(f"The Xenium patch shape must be two-dimensional, got {patch_shape}.")
469    if label_channel not in LABEL_CHANNELS:
470        raise ValueError(f"'{label_channel}' is not a valid label channel. Choose from {list(LABEL_CHANNELS)}.")
471    if raw_channel not in ("dapi", "stack"):
472        raise ValueError(f"'{raw_channel}' is not a valid raw channel. Choose from ['dapi', 'stack'].")
473
474    volume_paths = get_xenium_paths(path, sample, download, with_stains=raw_channel == "stack")
475
476    kwargs = util.update_kwargs(kwargs, "with_channels", raw_channel == "stack")
477    kwargs, _ = util.add_instance_label_transform(
478        kwargs, add_binary_target=True, offsets=offsets, boundaries=boundaries, binary=binary,
479    )
480    kwargs = util.ensure_transforms(ndim=2, **kwargs)
481
482    return torch_em.default_segmentation_dataset(
483        raw_paths=volume_paths,
484        raw_key=f"raw/{raw_channel}",
485        label_paths=volume_paths,
486        label_key=f"labels/{label_channel}",
487        patch_shape=patch_shape,
488        ndim=2,
489        **kwargs,
490    )
491
492
493def get_xenium_loader(
494    path: Union[os.PathLike, str],
495    batch_size: int,
496    patch_shape: Tuple[int, int],
497    sample: Optional[Union[str, Sequence[str]]] = None,
498    label_channel: Literal["nuclei", "cells"] = "nuclei",
499    raw_channel: Literal["dapi", "stack"] = "dapi",
500    offsets: Optional[List[List[int]]] = None,
501    boundaries: bool = False,
502    binary: bool = False,
503    download: bool = False,
504    **kwargs,
505) -> DataLoader:
506    """Get the Xenium dataloader for cell and nucleus segmentation.
507
508    Args:
509        path: Filepath to a folder where the downloaded data will be saved.
510        batch_size: The batch size for training.
511        patch_shape: The 2D patch shape to use for training.
512        sample: The sample or samples to use. Defaults to all of them.
513        label_channel: The masks to use as target. Either 'nuclei' or 'cells'.
514        raw_channel: The images to use as input. Either 'dapi' or 'stack'.
515        offsets: Offset values for affinity computation used as target.
516        boundaries: Whether to compute boundaries as the target.
517        binary: Whether to use a binary segmentation target.
518        download: Whether to download the data if it is not present.
519        kwargs: Additional keyword arguments for `torch_em.default_segmentation_dataset` or the PyTorch DataLoader.
520
521    Returns:
522        The DataLoader.
523    """
524    ds_kwargs, loader_kwargs = util.split_kwargs(torch_em.default_segmentation_dataset, **kwargs)
525    dataset = get_xenium_dataset(
526        path=path,
527        patch_shape=patch_shape,
528        sample=sample,
529        label_channel=label_channel,
530        raw_channel=raw_channel,
531        offsets=offsets,
532        boundaries=boundaries,
533        binary=binary,
534        download=download,
535        **ds_kwargs,
536    )
537    return torch_em.get_data_loader(dataset, batch_size=batch_size, **loader_kwargs)