torch_em.data.datasets.util

  1import os
  2import hashlib
  3import inspect
  4import zipfile
  5import requests
  6from tqdm import tqdm
  7from warnings import warn
  8from subprocess import run
  9from xml.dom import minidom
 10from packaging import version
 11from shutil import copyfileobj, which
 12
 13from typing import Optional, Tuple, Literal, List, Dict, Union, Callable
 14
 15import numpy as np
 16from skimage.draw import polygon
 17
 18import torch
 19
 20import torch_em
 21from torch_em.transform import get_raw_transform
 22from torch_em.transform.generic import ResizeLongestSideInputs, Compose
 23
 24try:
 25    import gdown
 26except ImportError:
 27    gdown = None
 28
 29try:
 30    from tcia_utils import nbia
 31except ModuleNotFoundError:
 32    nbia = None
 33
 34try:
 35    from cryoet_data_portal import Client, Dataset
 36except ImportError:
 37    Client, Dataset = None, None
 38
 39try:
 40    import synapseclient
 41    import synapseutils
 42except ImportError:
 43    synapseclient, synapseutils = None, None
 44
 45
 46# fsspec's HTTP filesystem (aiohttp) ignores the http(s)_proxy environment variables unless 'trust_env' is set,
 47# so remote data behind a proxy needs these storage options.
 48HTTP_STORAGE_OPTIONS = {"client_kwargs": {"trust_env": True}}
 49
 50
 51BIOIMAGEIO_IDS = {
 52    "covid_if": "ilastik/covid_if_training_data",
 53    "cremi": "ilastik/cremi_training_data",
 54    "dsb": "ilastik/stardist_dsb_training_data",
 55    "hpa": "",  # not on bioimageio yet
 56    "isbi2012": "ilastik/isbi2012_neuron_segmentation_challenge",
 57    "kasthuri": "",  # not on bioimageio yet:
 58    "livecell": "ilastik/livecell_dataset",
 59    "lucchi": "",  # not on bioimageio yet:
 60    "mitoem": "ilastik/mitoem_segmentation_challenge",
 61    "monuseg": "deepimagej/monuseg_digital_pathology_miccai2018",
 62    "ovules": "",  # not on bioimageio yet
 63    "plantseg_root": "ilastik/plantseg_root",
 64    "plantseg_ovules": "ilastik/plantseg_ovules",
 65    "platynereis": "ilastik/platynereis_em_training_data",
 66    "snemi": "",  # not on bioimagegio yet
 67    "uro_cell": "",  # not on bioimageio yet: https://doi.org/10.1016/j.compbiomed.2020.103693
 68    "vnc": "ilastik/vnc",
 69}
 70"""@private
 71"""
 72
 73
 74def get_bioimageio_dataset_id(dataset_name):
 75    """@private
 76    """
 77    assert dataset_name in BIOIMAGEIO_IDS
 78    return BIOIMAGEIO_IDS[dataset_name]
 79
 80
 81def get_checksum(filename: str) -> str:
 82    """Get the SHA256 checksum of a file.
 83
 84    Args:
 85        filename: The filepath.
 86
 87    Returns:
 88        The checksum.
 89    """
 90    # The file is hashed in chunks, so that datasets with multi-GB archives do not run out of memory.
 91    hasher = hashlib.sha256()
 92    with open(filename, "rb") as f:
 93        for chunk in iter(lambda: f.read(64 * 1024 * 1024), b""):
 94            hasher.update(chunk)
 95    return hasher.hexdigest()
 96
 97
 98def _check_checksum(path, checksum):
 99    if checksum is not None:
100        this_checksum = get_checksum(path)
101        if this_checksum != checksum:
102            raise RuntimeError(
103                "The checksum of the download does not match the expected checksum."
104                f"Expected: {checksum}, got: {this_checksum}"
105            )
106        print("Download successful and checksums agree.")
107    else:
108        warn("The file was downloaded, but no checksum was provided, so the file may be corrupted.")
109
110
111# this needs to be extended to support download from s3 via boto,
112# if we get a resource that is available via s3 without support for http
113def download_source(path: str, url: str, download: bool, checksum: Optional[str] = None, verify: bool = True) -> None:
114    """Download data via https.
115
116    Args:
117        path: The path for saving the data.
118        url: The url of the data.
119        download: Whether to download the data if it is not saved at `path` yet.
120        checksum: The expected checksum of the data.
121        verify: Whether to verify the https address.
122    """
123    if os.path.exists(path):
124        return
125    if not download:
126        raise RuntimeError(f"Cannot find the data at {path}, but download was set to False")
127
128    # The data is downloaded to a temporary path and only moved to `path` once it is complete and verified.
129    # Otherwise an interrupted download would be mistaken for a complete one by the check above.
130    tmp_path = f"{path}.incomplete"
131    with requests.get(url, stream=True, allow_redirects=True, verify=verify) as r:
132        r.raise_for_status()  # check for error
133        # Compute checksums on the file content rather than its HTTP transfer encoding.
134        r.raw.decode_content = True
135        file_size = int(r.headers.get("Content-Length", 0))
136        desc = f"Download {url} to {path}"
137        if file_size == 0:
138            desc += " (unknown file size)"
139        with tqdm.wrapattr(r.raw, "read", total=file_size, desc=desc) as r_raw, open(tmp_path, "wb") as f:
140            copyfileobj(r_raw, f)
141
142    _check_checksum(tmp_path, checksum)
143    os.replace(tmp_path, path)
144
145
146def download_source_gdrive(
147    path: str,
148    url: str,
149    download: bool,
150    checksum: Optional[str] = None,
151    download_type: Literal["zip", "folder"] = "zip",
152    expected_samples: int = 10000,
153    quiet: bool = True,
154) -> None:
155    """Download data from google drive.
156
157    Args:
158        path: The path for saving the data.
159        url: The url of the data.
160        download: Whether to download the data if it is not saved at `path` yet.
161        checksum: The expected checksum of the data.
162        download_type: The download type, either 'zip' or 'folder'.
163        expected_samples: The maximal number of samples in the folder.
164        quiet: Whether to download quietly.
165    """
166    if os.path.exists(path):
167        return
168
169    if not download:
170        raise RuntimeError(f"Cannot find the data at {path}, but download was set to False")
171
172    if gdown is None:
173        raise RuntimeError(
174            "Need gdown library to download data from google drive. "
175            "Please install gdown: 'conda install -c conda-forge gdown==4.6.3'."
176        )
177
178    print("Downloading the files. Might take a few minutes...")
179
180    if download_type == "zip":
181        gdown.download(url, path, quiet=quiet)
182        _check_checksum(path, checksum)
183    elif download_type == "folder":
184        assert version.parse(gdown.__version__) == version.parse("4.6.3"), "Please install 'gdown==4.6.3'."
185        gdown.download_folder.__globals__["MAX_NUMBER_FILES"] = expected_samples
186        gdown.download_folder(url=url, output=path, quiet=quiet, remaining_ok=True)
187    else:
188        raise ValueError("`download_path` argument expects either `zip`/`folder`")
189
190    print("Download completed.")
191
192
193def download_source_empiar(path: str, access_id: str, download: bool) -> str:
194    """Download data from EMPIAR.
195
196    Requires the ascp command from the aspera CLI.
197
198    Args:
199        path: The path for saving the data.
200        access_id: The EMPIAR accession id of the data to download.
201        download: Whether to download the data if it is not saved at `path` yet.
202
203    Returns:
204        The path to the downloaded data.
205    """
206    download_path = os.path.join(path, access_id)
207
208    if os.path.exists(download_path):
209        return download_path
210    if not download:
211        raise RuntimeError(f"Cannot find the data at {path}, but download was set to False")
212
213    if which("ascp") is None:
214        raise RuntimeError(
215            "Need aspera-cli to download data from empiar. You can install it via 'conda install -c hcc aspera-cli'."
216        )
217
218    key_file = os.path.expanduser("~/.aspera/cli/etc/asperaweb_id_dsa.openssh")
219    if not os.path.exists(key_file):
220        conda_root = os.environ["CONDA_PREFIX"]
221        key_file = os.path.join(conda_root, "etc/asperaweb_id_dsa.openssh")
222
223    if not os.path.exists(key_file):
224        raise RuntimeError("Could not find the aspera ssh keyfile")
225
226    cmd = ["ascp", "-QT", "-l", "200M", "-P33001", "-i", key_file, f"emp_ext2@fasp.ebi.ac.uk:/{access_id}", path]
227    run(cmd)
228
229    return download_path
230
231
232def download_source_kaggle(path: str, dataset_name: str, download: bool, competition: bool = False):
233    """Download data from Kaggle.
234
235    Requires the Kaggle API.
236
237    Args:
238        path: The path for saving the data.
239        dataset_name: The name of the dataset to download.
240        download: Whether to download the data if it is not saved at `path` yet.
241        competition: Whether this data is from a competition and requires the kaggle.competition API.
242    """
243    if not download:
244        raise RuntimeError(f"Cannot find the data at {path}, but download was set to False.")
245
246    try:
247        from kaggle.api.kaggle_api_extended import KaggleApi
248    except ModuleNotFoundError:
249        msg = "Please install the Kaggle API. You can do this using 'pip install kaggle'. "
250        msg += "After you have installed kaggle, you would need an API token. "
251        msg += "Follow the instructions at https://www.kaggle.com/docs/api."
252        raise ModuleNotFoundError(msg)
253
254    api = KaggleApi()
255    api.authenticate()
256
257    if competition:
258        api.competition_download_files(competition=dataset_name, path=path, quiet=False)
259    else:
260        api.dataset_download_files(dataset=dataset_name, path=path, quiet=False)
261
262
263NBIA_API_URL = "https://services.cancerimagingarchive.net/nbia-api/services/v1/"
264
265
266def _download_tcia_series_with_rest(series_uids, dst, csv_filename):
267    """Download DICOM series from TCIA via the NBIA REST API.
268
269    This is the fallback for missing 'tcia_utils'. It mimics the on-disk layout of 'nbia.downloadSeries':
270    each series is extracted to '<dst>/<SeriesInstanceUID>/' and the series metadata are written to
271    '<csv_filename>.csv' (including the 'Series UID', 'Subject ID' and 'Modality' columns).
272    The metadata of each series is cached in '<csv_filename>.partial.json' while the download is running,
273    so that an interrupted download does not have to query the metadata of all series again.
274    See https://wiki.cancerimagingarchive.net/x/fILTB for the API documentation.
275    """
276    import csv
277    import json
278    import time
279    import tempfile
280
281    def get_with_retries(endpoint, n_retries=5, **kwargs):
282        # The NBIA API occasionally returns server errors, so the requests are retried.
283        for attempt in range(n_retries):
284            response = requests.get(NBIA_API_URL + endpoint, **kwargs)
285            if response.status_code < 500 or attempt == n_retries - 1:
286                response.raise_for_status()
287                return response
288            response.close()
289            time.sleep(10 * (attempt + 1))
290
291    os.makedirs(dst, exist_ok=True)
292    cache_path = f"{csv_filename}.partial.json"
293    cache = {}
294    if os.path.exists(cache_path):
295        with open(cache_path, "r") as f:
296            cache = json.load(f)
297
298    metadata = []
299    for i, uid in enumerate(tqdm(series_uids, desc=f"Download {len(series_uids)} series from TCIA to {dst}")):
300        series_dir = os.path.join(dst, uid)
301        if uid in cache and os.path.exists(series_dir):
302            metadata.append(cache[uid])
303            continue
304
305        response = get_with_retries("getSeriesMetaData", params={"SeriesInstanceUID": uid})
306        # A small number of series have no indexed metadata and the endpoint returns an empty body for
307        # them, even though the image data itself downloads without issue. A transient proxy or gateway
308        # error can also return a non-JSON body with a 200 status, which is handled the same way.
309        try:
310            rows = response.json() if response.content else []
311        except ValueError:
312            rows = []
313        metadata.append(rows[0] if rows else {"Series UID": uid})
314        cache[uid] = metadata[-1]
315        if i % 50 == 0:
316            with open(cache_path, "w") as f:
317                json.dump(cache, f)
318
319        if os.path.exists(series_dir):  # This series has been downloaded already.
320            continue
321
322        # The series is downloaded as a zip archive, which is extracted to a temporary folder
323        # and only moved to the final location once it is complete.
324        try:
325            with tempfile.TemporaryDirectory(dir=dst) as tmp_dir:
326                zip_path = os.path.join(tmp_dir, "series.zip")
327                with get_with_retries("getImage", params={"SeriesInstanceUID": uid}, stream=True) as r:
328                    with open(zip_path, "wb") as f:
329                        copyfileobj(r.raw, f)
330                tmp_series_dir = os.path.join(tmp_dir, "series")
331                unzip(zip_path, tmp_series_dir)
332                os.rename(tmp_series_dir, series_dir)
333        except FileExistsError:
334            # A resumed download may have already extracted this series: on some network filesystems
335            # the 'os.path.exists' check above can be stale, so this is not caught earlier.
336            pass
337        except requests.exceptions.HTTPError as e:
338            # A small number of series are indexed but no longer resolvable through this endpoint (e.g. a
339            # stale UID), which should not abort downloading the rest of a possibly multi-hour batch.
340            if e.response is not None and e.response.status_code < 500:
341                tqdm.write(f"Skipping {uid}, which failed to download: {e}")
342            else:
343                raise
344
345    # The metadata keys differ between series (e.g. 'Series Date' is only reported for some), so the header
346    # has to be the union of all keys.
347    fieldnames = list(dict.fromkeys(key for row in metadata for key in row))
348    csv_path = f"{csv_filename}.csv"
349    with open(csv_path, "w", newline="") as f:
350        writer = csv.DictWriter(f, fieldnames=fieldnames)
351        writer.writeheader()
352        writer.writerows(metadata)
353
354    if os.path.exists(cache_path):  # The download is complete, so the metadata cache is not needed anymore.
355        os.remove(cache_path)
356
357
358def _download_tcia_manifest_with_rest(manifest_path, dst, csv_filename):
359    """Download all series listed in a TCIA manifest via the NBIA REST API."""
360    with open(manifest_path, "r") as f:
361        lines = [line.strip() for line in f.readlines()]
362    series_uids = lines[lines.index("ListOfSeriesToDownload=") + 1:]
363    series_uids = [uid for uid in series_uids if uid]
364    _download_tcia_series_with_rest(series_uids, dst, csv_filename)
365
366
367def download_tcia_series(series_uids: List[str], dst: str, csv_filename: str) -> str:
368    """Download individual DICOM series from TCIA by their series instance UIDs.
369
370    Uses the tcia_utils python package if it is installed and falls back to the NBIA REST API otherwise.
371    Each series is stored in '<dst>/<SeriesInstanceUID>/', series that exist there already are skipped.
372
373    Args:
374        series_uids: The UIDs of the series to download.
375        dst: The folder for saving the DICOM series.
376        csv_filename: The path for saving the series metadata (without the '.csv' extension).
377
378    Returns:
379        The path to the csv file with the series metadata.
380    """
381    if nbia is None:
382        _download_tcia_series_with_rest(series_uids, dst, csv_filename)
383    else:
384        nbia.downloadSeries(series_data=series_uids, input_type="list", path=dst, csv_filename=csv_filename)
385    return f"{csv_filename}.csv"
386
387
388def download_source_tcia(path, url, dst, csv_filename, download):
389    """Download data from TCIA.
390
391    Uses the tcia_utils python package if it is installed and falls back to the NBIA REST API otherwise.
392
393    Args:
394        path: The path for saving the manifest file. If `url` is None, this must point to an existing manifest,
395            e.g. one that was written by the caller to download only a subset of the series of a collection.
396        url: The URL to the TCIA manifest of the dataset. Set to None to use the manifest at `path`.
397        dst: The folder for saving the DICOM series. Each series is stored in a sub-folder named after its UID.
398        csv_filename: The path for saving the series metadata (without the '.csv' extension).
399        download: Whether to download the data if it is not saved at `path` yet.
400    """
401    if not download:
402        raise RuntimeError(f"Cannot find the data at {path}, but download was set to False.")
403
404    if url is None:
405        assert os.path.exists(path), f"The manifest {path} does not exist."
406    else:
407        assert url.endswith(".tcia"), f"{url} is not a TCIA Manifest."
408        # Downloads the manifest file from the collection page.
409        manifest = requests.get(url=url)
410        manifest.raise_for_status()
411        with open(path, "wb") as f:
412            f.write(manifest.content)
413
414    # This part extracts the UIDs from the manifests and downloads them.
415    if nbia is None:
416        _download_tcia_manifest_with_rest(path, dst, csv_filename)
417    else:
418        nbia.downloadSeries(series_data=path, input_type="manifest", path=dst, csv_filename=csv_filename)
419
420
421def download_source_synapse(path: str, entity: str, download: bool) -> None:
422    """Download data from synapse.
423
424    Requires the synapseclient python library.
425
426    Args:
427        path: The path for saving the data.
428        entity: The name of the data to download from synapse.
429        download: Whether to download the data if it is not saved at `path` yet.
430    """
431    if not download:
432        raise RuntimeError(f"Cannot find the data at {path}, but download was set to False.")
433
434    if synapseclient is None:
435        raise RuntimeError(
436            "You must install 'synapseclient' to download files from 'synapse'. "
437            "Remember to create an account and generate an authentication code for your account. "
438            "Please follow the documentation for details on creating the '~/.synapseConfig' file here: "
439            "https://python-docs.synapse.org/tutorials/authentication/."
440        )
441
442    assert entity.startswith("syn"), "The entity name does not look as expected. It should be something like 'syn123'."
443
444    # Download all files in the folder.
445    syn = synapseclient.Synapse()
446    syn.login()  # Since we do not pass any credentials here, it fetches all details from '~/.synapseConfig'.
447    synapseutils.syncFromSynapse(syn=syn, entity=entity, path=path)
448
449
450def update_kwargs(kwargs, key, value, msg=None):
451    """@private
452    """
453    if key in kwargs:
454        msg = f"{key} will be over-ridden in loader kwargs." if msg is None else msg
455        warn(msg)
456    kwargs[key] = value
457    return kwargs
458
459
460def unzip_tarfile(tar_path: str, dst: str, remove: bool = True) -> None:
461    """Unpack a tar archive.
462
463    Args:
464        tar_path: Path to the tar file.
465        dst: Where to unpack the archive.
466        remove: Whether to remove the tar file after unpacking.
467    """
468    import tarfile
469
470    if tar_path.endswith(".tar.gz") or tar_path.endswith(".tgz"):
471        access_mode = "r:gz"
472    elif tar_path.endswith(".tar"):
473        access_mode = "r:"
474    else:
475        raise ValueError(f"The provided file isn't a supported archive to unpack. Please check the file: {tar_path}.")
476
477    tar = tarfile.open(tar_path, access_mode)
478    tar.extractall(dst)
479    tar.close()
480
481    if remove:
482        os.remove(tar_path)
483
484
485def unzip_rarfile(rar_path: str, dst: str, remove: bool = True, use_rarfile: bool = True) -> None:
486    """Unpack a rar archive.
487
488    Args:
489        rar_path: Path to the rar file.
490        dst: Where to unpack the archive.
491        remove: Whether to remove the tar file after unpacking.
492        use_rarfile: Whether to use the rarfile library or aspose.zip.
493    """
494    def _extract_with_rarfile():
495        import rarfile
496        with rarfile.RarFile(rar_path) as archive:
497            archive.extractall(path=dst)
498
499    def _extract_with_aspose():
500        import aspose.zip as az
501        with az.rar.RarArchive(rar_path) as archive:
502            archive.extract_to_directory(dst)
503
504    def _extract_with_7z():
505        if which("7z") is None:
506            raise RuntimeError("The 'p7zip' CLI is not available.")
507        run(["7z", "x", f"-o{dst}", "-y", rar_path], check=True)
508
509    extractors = [
510        ('rarfile', _extract_with_rarfile), ('aspose.zip', _extract_with_aspose), ('7z', _extract_with_7z),
511    ] if use_rarfile else [('aspose.zip', _extract_with_aspose), ('7z', _extract_with_7z)]
512
513    errors = []
514    for name, extractor in extractors:
515        try:
516            extractor()
517            break
518        except Exception as err:
519            errors.append((name, err))
520            if len(errors) < len(extractors):
521                next_name = extractors[len(errors)][0]
522                warn(f"Extraction with '{name}' failed for {rar_path} ({err}). Falling back to '{next_name}'.")
523    else:
524        backends = ', '.join(f"'{name}'" for name, _ in extractors)
525        raise RuntimeError(
526            f"Failed to extract rar archive {rar_path} with {backends}. "
527            "Please ensure one of the supported backends is installed and can read this archive."
528        ) from errors[-1][1]
529
530    if remove:
531        os.remove(rar_path)
532
533
534def unzip(zip_path: str, dst: str, remove: bool = True) -> None:
535    """Unpack a zip archive.
536
537    Args:
538        zip_path: Path to the zip file.
539        dst: Where to unpack the archive.
540        remove: Whether to remove the tar file after unpacking.
541    """
542    with zipfile.ZipFile(zip_path, "r") as f:
543        f.extractall(dst)
544    if remove:
545        os.remove(zip_path)
546
547
548def unzip_7z(path_7z: str, dst: str, remove: bool = True) -> None:
549    """Unpack a 7z archive.
550
551    Args:
552        path_7z: Path to the 7z file.
553        dst: Where to unpack the archive.
554        remove: Whether to remove the 7z file after unpacking.
555    """
556    if which("7z") is None:
557        raise RuntimeError("Need the 'p7zip' CLI to extract 7z archives. You can install it via 'conda install -c conda-forge p7zip'.")  # noqa
558
559    run(["7z", "x", f"-o{dst}", "-y", path_7z])
560
561    if remove:
562        os.remove(path_7z)
563
564
565def split_kwargs(function, **kwargs):
566    """@private
567    """
568    function_parameters = inspect.signature(function).parameters
569    parameter_names = list(function_parameters.keys())
570    other_kwargs = {k: v for k, v in kwargs.items() if k not in parameter_names}
571    kwargs = {k: v for k, v in kwargs.items() if k in parameter_names}
572    return kwargs, other_kwargs
573
574
575# this adds the default transforms for 'raw_transform' and 'transform'
576# in case these were not specified in the kwargs
577# this is NOT necessary if 'default_segmentation_dataset' is used, only if a dataset class
578# is used directly, e.g. in the LiveCell Loader
579def ensure_transforms(ndim, **kwargs):
580    """@private
581    """
582    if "raw_transform" not in kwargs:
583        kwargs = update_kwargs(kwargs, "raw_transform", torch_em.transform.get_raw_transform())
584    if "transform" not in kwargs:
585        kwargs = update_kwargs(kwargs, "transform", torch_em.transform.get_augmentations(ndim=ndim))
586    return kwargs
587
588
589def add_instance_label_transform(
590    kwargs, add_binary_target, label_dtype=None, binary=False, boundaries=False, offsets=None, binary_is_exclusive=True,
591):
592    """@private
593    """
594    if binary_is_exclusive:
595        assert sum((offsets is not None, boundaries, binary)) <= 1
596    else:
597        assert sum((offsets is not None, boundaries)) <= 1
598    if offsets is not None:
599        label_transform2 = torch_em.transform.label.AffinityTransform(offsets=offsets,
600                                                                      add_binary_target=add_binary_target,
601                                                                      add_mask=True)
602        msg = "Offsets are passed, but 'label_transform2' is in the kwargs. It will be over-ridden."
603        kwargs = update_kwargs(kwargs, "label_transform2", label_transform2, msg=msg)
604        label_dtype = torch.float32
605    elif boundaries:
606        label_transform = torch_em.transform.label.BoundaryTransform(add_binary_target=add_binary_target)
607        msg = "Boundaries is set to true, but 'label_transform' is in the kwargs. It will be over-ridden."
608        kwargs = update_kwargs(kwargs, "label_transform", label_transform, msg=msg)
609        label_dtype = torch.float32
610    elif binary:
611        label_transform = torch_em.transform.label.labels_to_binary
612        msg = "Binary is set to true, but 'label_transform' is in the kwargs. It will be over-ridden."
613        kwargs = update_kwargs(kwargs, "label_transform", label_transform, msg=msg)
614        label_dtype = torch.float32
615    return kwargs, label_dtype
616
617
618def update_kwargs_for_resize_trafo(kwargs, patch_shape, resize_inputs, resize_kwargs=None, ensure_rgb=None):
619    """@private
620    """
621    # Checks for raw_transform and label_transform incoming values.
622    # If yes, it will automatically merge these two transforms to apply them together.
623    if resize_inputs:
624        assert isinstance(resize_kwargs, dict)
625
626        target_shape = resize_kwargs.get("patch_shape")
627        if len(resize_kwargs["patch_shape"]) == 3:
628            # we only need the XY dimensions to reshape the inputs along them.
629            target_shape = target_shape[1:]
630            # we provide the Z dimension value to return the desired number of slices and not the whole volume
631            kwargs["z_ext"] = resize_kwargs["patch_shape"][0]
632
633        raw_trafo = ResizeLongestSideInputs(target_shape=target_shape, is_rgb=resize_kwargs["is_rgb"])
634        label_trafo = ResizeLongestSideInputs(target_shape=target_shape, is_label=True)
635
636        # The patch shape provided to the dataset. Here, "None" means that the entire volume will be loaded.
637        patch_shape = None
638
639    if ensure_rgb is None:
640        raw_trafos = []
641    else:
642        assert not isinstance(ensure_rgb, bool), "'ensure_rgb' is expected to be a function."
643        raw_trafos = [ensure_rgb]
644
645    if "raw_transform" in kwargs:
646        raw_trafos.extend([raw_trafo, kwargs["raw_transform"]])
647    else:
648        raw_trafos.extend([raw_trafo, get_raw_transform()])
649
650    kwargs["raw_transform"] = Compose(*raw_trafos, is_multi_tensor=False)
651
652    if "label_transform" in kwargs:
653        trafo = Compose(label_trafo, kwargs["label_transform"], is_multi_tensor=False)
654        kwargs["label_transform"] = trafo
655    else:
656        kwargs["label_transform"] = label_trafo
657
658    return kwargs, patch_shape
659
660
661def generate_labeled_array_from_xml(shape: Tuple[int, ...], xml_file: str) -> np.ndarray:
662    """Generate a label mask from a contour defined in a xml annotation file.
663
664    Function taken from: https://github.com/rshwndsz/hover-net/blob/master/lightning_hovernet.ipynb
665
666    Args:
667        shape: The image shape.
668        xml_file: The path to the xml file with contour annotations.
669
670    Returns:
671        The label mask.
672    """
673    # DOM object created by the minidom parser
674    xDoc = minidom.parse(xml_file)
675
676    # List of all Region tags
677    regions = xDoc.getElementsByTagName('Region')
678
679    # List which will store the vertices for each region
680    xy = []
681    for region in regions:
682        # Loading all the vertices in the region
683        vertices = region.getElementsByTagName('Vertex')
684
685        # The vertices of a region will be stored in a array
686        vw = np.zeros((len(vertices), 2))
687
688        for index, vertex in enumerate(vertices):
689            # Storing the values of x and y coordinate after conversion
690            vw[index][0] = float(vertex.getAttribute('X'))
691            vw[index][1] = float(vertex.getAttribute('Y'))
692
693        # Append the vertices of a region
694        xy.append(np.int32(vw))
695
696    # Creating a completely black image
697    mask = np.zeros(shape, np.uint32)  # Integer instance ids; float labels break connected-component ops.
698
699    # Start the instance ids at 1: id 0 is background, so enumerating from 0 silently drops the first region.
700    for i, contour in enumerate(xy, start=1):
701        r, c = polygon(np.array(contour)[:, 1], np.array(contour)[:, 0], shape=shape)
702        mask[r, c] = i
703    return mask
704
705
706def load_dicom_series(series_dir: str) -> Tuple[np.ndarray, Dict[str, np.ndarray]]:
707    """Stack a single-frame DICOM image series (CT, MR, PET) into a volume with axes (z, y, x).
708
709    The slices are sorted by their position along the slice normal (the cross product of the row and column
710    direction in 'ImageOrientationPatient'), so the volume is stacked consistently for any acquisition plane.
711    'RescaleSlope' and 'RescaleIntercept' are applied per slice, i.e. CT volumes are returned in Hounsfield units.
712
713    NOTE: This requires the pydicom python package.
714
715    Args:
716        series_dir: The folder with the DICOM files of the series.
717
718    Returns:
719        The volume with axes (z, y, x) as float32.
720        The geometry of the volume, which is needed by `rasterize_rtstruct`. A dictionary with the keys
721        'origin' (the 'ImagePositionPatient' of each slice, n_slices x 3), 'row_direction' and 'column_direction'
722        (the unit vectors along which the column and the row index increase, from 'ImageOrientationPatient'),
723        'spacing' (the row and column spacing from 'PixelSpacing') and 'sop_uids' (the 'SOPInstanceUID' per slice).
724    """
725    import pydicom
726
727    dcm_paths = [os.path.join(series_dir, fname) for fname in sorted(os.listdir(series_dir)) if fname.endswith(".dcm")]
728    slices = [pydicom.dcmread(dcm_path) for dcm_path in dcm_paths]
729
730    row_direction = np.array([float(v) for v in slices[0].ImageOrientationPatient[:3]])
731    column_direction = np.array([float(v) for v in slices[0].ImageOrientationPatient[3:]])
732    normal = np.cross(row_direction, column_direction)
733    slices.sort(key=lambda dcm: np.dot([float(v) for v in dcm.ImagePositionPatient], normal))
734
735    volume = []
736    for dcm in slices:
737        frame = dcm.pixel_array.astype("float32")
738        volume.append(frame * float(dcm.get("RescaleSlope", 1.0)) + float(dcm.get("RescaleIntercept", 0.0)))
739    volume = np.stack(volume)
740
741    geometry = {
742        "origin": np.array([[float(v) for v in dcm.ImagePositionPatient] for dcm in slices]),
743        "row_direction": row_direction,
744        "column_direction": column_direction,
745        "spacing": np.array([float(v) for v in slices[0].PixelSpacing]),
746        "sop_uids": np.array([str(dcm.SOPInstanceUID) for dcm in slices]),
747    }
748    return volume, geometry
749
750
751def rasterize_rtstruct(
752    rtstruct_path: str,
753    geometry: Dict[str, np.ndarray],
754    shape: Tuple[int, int, int],
755    roi_labels: Union[Dict[str, int], Callable[[int, str], Optional[int]]],
756) -> np.ndarray:
757    """Rasterize the contours of a DICOM RTSTRUCT file onto the voxel grid of the referenced image series.
758
759    Each 'CLOSED_PLANAR' contour is assigned to the slice it references (via 'ReferencedSOPInstanceUID', with a
760    fallback to the closest slice along the slice normal). Its points are projected onto the row and column
761    direction of that slice to obtain pixel coordinates and the polygon is filled with `skimage.draw.polygon`.
762    Multiple contours of the same ROI on the same slice are combined with XOR, so that inner contours form holes.
763    Where different ROIs overlap, the ROI with the lower label id takes precedence.
764
765    NOTE: This requires the pydicom python package.
766
767    Args:
768        rtstruct_path: The path to the RTSTRUCT DICOM file.
769        geometry: The geometry of the referenced image series, as returned by `load_dicom_series`.
770        shape: The shape of the image volume (z, y, x).
771        roi_labels: The mapping from ROIs to label ids. Either a dictionary that maps the ROI names to label ids,
772            or a function that maps the ROI number and ROI name to a label id. ROIs that are not in the dictionary
773            or for which the function returns None are ignored.
774
775    Returns:
776        The label volume (uint8) with axes (z, y, x).
777    """
778    import pydicom
779
780    rtstruct = pydicom.dcmread(rtstruct_path)
781    roi_names = {int(roi.ROINumber): str(roi.ROIName) for roi in rtstruct.StructureSetROISequence}
782    slice_ids = {uid: z for z, uid in enumerate(geometry["sop_uids"])}
783    normal = np.cross(geometry["row_direction"], geometry["column_direction"])
784    slice_positions = geometry["origin"] @ normal
785
786    masks = {}
787    for roi_contour in rtstruct.ROIContourSequence:
788        roi_number = int(roi_contour.ReferencedROINumber)
789        roi_name = roi_names[roi_number]
790        if isinstance(roi_labels, dict):
791            label_id = roi_labels.get(roi_name)
792        else:
793            label_id = roi_labels(roi_number, roi_name)
794        if label_id is None:
795            continue
796
797        mask = masks.setdefault(label_id, np.zeros(shape, dtype="bool"))
798        for contour in roi_contour.get("ContourSequence", []):
799            if contour.ContourGeometricType != "CLOSED_PLANAR":
800                continue
801            points = np.array([float(v) for v in contour.ContourData]).reshape(-1, 3)
802            if len(points) < 3:
803                continue
804
805            z = None
806            for ref in contour.get("ContourImageSequence", []):
807                z = slice_ids.get(str(ref.ReferencedSOPInstanceUID))
808            if z is None:
809                z = int(np.argmin(np.abs(slice_positions - np.dot(points[0], normal))))
810
811            offsets = points - geometry["origin"][z]
812            cols = offsets @ geometry["row_direction"] / geometry["spacing"][1]
813            rows = offsets @ geometry["column_direction"] / geometry["spacing"][0]
814            rr, cc = polygon(rows, cols, shape=shape[1:])
815            mask[z, rr, cc] = ~mask[z, rr, cc]
816
817    labels = np.zeros(shape, dtype="uint8")
818    for label_id in sorted(masks, reverse=True):
819        labels[masks[label_id]] = label_id
820    return labels
821
822
823# This function could be extended to convert WSIs (or modalities with multiple resolutions).
824def convert_svs_to_array(
825    path: str, location: Tuple[int, int] = (0, 0), level: int = 0, img_size: Tuple[int, int] = None,
826) -> np.ndarray:
827    """Convert a .svs file for WSI imagging to a numpy array.
828
829    Requires the tiffslide python library.
830    The function can load multi-resolution images. You can specify the resolution level via `level`.
831
832    Args:
833        path: File path ath to the svs file.
834        location: Pixel location (x, y) in level 0 of the image.
835        level: Target level used to read the image.
836        img_size: Size of the image. If None, the shape of the image at `level` is used.
837
838    Returns:
839        The image as numpy array.
840    """
841    assert path.endswith(".svs"), f"The provided file ({path}) isn't in svs format"
842
843    try:
844        from tiffslide import TiffSlide
845    except ImportError:
846        # svs is a pyramidal TIFF variant, so tifffile can read it without the tiffslide dependency.
847        import tifffile
848        with tifffile.TiffFile(path) as f:
849            image = f.series[0].levels[level].asarray()
850        x, y = location
851        if img_size is not None:
852            image = image[y:y + img_size[1], x:x + img_size[0]]
853        else:
854            image = image[y:, x:]
855        return image
856
857    _slide = TiffSlide(path)
858    if img_size is None:
859        img_size = _slide.level_dimensions[0]
860    return _slide.read_region(location=location, level=level, size=img_size, as_array=True)
861
862
863def download_from_cryo_et_portal(path: str, dataset_id: int, download: bool) -> str:
864    """Download data from the CryoET Data Portal.
865
866    Requires the cryoet-data-portal python library.
867
868    Args:
869        path: The path for saving the data.
870        dataset_id: The id of the data to download from the portal.
871        download: Whether to download the data if it is not saved at `path` yet.
872
873    Returns:
874        The file path to the downloaded data.
875    """
876    if Client is None or Dataset is None:
877        raise RuntimeError("Please install CryoETDataPortal via 'pip install cryoet-data-portal'")
878
879    output_path = os.path.join(path, str(dataset_id))
880    if os.path.exists(output_path):
881        return output_path
882
883    if not download:
884        raise RuntimeError(f"Cannot find the data at {path}, but download was set to False.")
885
886    client = Client()
887    dataset = Dataset.get_by_id(client, dataset_id)
888    dataset.download_everything(dest_path=path)
889
890    return output_path
891
892
893CRYO_ET_PORTAL_BUCKET = "cryoet-data-portal-public"
894
895
896def _open_cryo_et_portal_zarr(s3_path):
897    import zarr
898    from zarr.storage import FsspecStore
899
900    store = FsspecStore.from_url(s3_path, storage_options={"anon": True}, read_only=True)
901    return zarr.open(store, mode="r")
902
903
904def stream_from_cryo_et_portal(
905    zarr_path: str,
906    dataset_id: int,
907    run: str,
908    voxel_spacing: str,
909    tomogram_folder: str,
910    annotations: Dict[str, str],
911    download: bool,
912) -> str:
913    """Stream a tomogram and its segmentation masks of a single run from the CryoET Data Portal.
914
915    In contrast to `download_from_cryo_et_portal` this only fetches the tomogram and the requested masks
916    (and not the frames and tilt series). The data is cached in a zarr group with the tomogram stored as 'raw'
917    and one array per annotation. Already cached arrays are not downloaded again.
918
919    Requires zarr (v3) and s3fs.
920
921    Args:
922        zarr_path: The filepath for the cached zarr group.
923        dataset_id: The id of the dataset on the portal.
924        run: The name of the run.
925        voxel_spacing: The voxel spacing folder of the run, e.g. 'VoxelSpacing13.480'.
926        tomogram_folder: The folder of the tomogram within 'Tomograms', e.g. '100'.
927        annotations: Mapping from the annotation name (e.g. 'nucleus') to the portal annotation folder (e.g. '108').
928        download: Whether to download the data if it is not cached at `zarr_path` yet.
929
930    Returns:
931        The filepath to the cached zarr group.
932    """
933    import zarr
934    from zarr.codecs import BloscCodec
935
936    root = zarr.open_group(zarr_path, mode="a")
937    if "raw" in root and all(name in root for name in annotations):
938        return zarr_path
939
940    if not download:
941        raise RuntimeError(f"Cannot find the data at {zarr_path}, but download was set to False.")
942
943    print(f"Streaming run {run} of dataset {dataset_id} from the CryoET Data Portal ...")
944    prefix = f"s3://{CRYO_ET_PORTAL_BUCKET}/{dataset_id}/{run}/Reconstructions/{voxel_spacing}"
945
946    def _make_array(name, data, shuffle):
947        arr = root.create_array(
948            name, shape=data.shape, chunks=tuple(min(128, s) for s in data.shape), dtype=data.dtype,
949            compressors=BloscCodec(cname="zstd", clevel=6, shuffle=shuffle),
950        )
951        arr[:] = data
952
953    raw = _open_cryo_et_portal_zarr(f"{prefix}/Tomograms/{tomogram_folder}/{run}.zarr/0")[:]
954    if "raw" not in root:
955        _make_array("raw", raw, shuffle="shuffle")
956
957    for name, folder in annotations.items():
958        if name in root:
959            continue
960        labels = _open_cryo_et_portal_zarr(f"{prefix}/Annotations/{folder}/{name}-1.0_segmentationmask.zarr/0")[:]
961        assert labels.shape == raw.shape, f"Shape mismatch for {run}, '{name}': raw {raw.shape} vs label {labels.shape}"
962        _make_array(name, labels, shuffle="bitshuffle")
963
964    root.attrs["dataset_id"] = dataset_id
965    root.attrs["run"] = run
966    return zarr_path