allow URLs for export-as-tiff (#1356)

* allow URLs for export-as-tiff

* also work with sharing_token

* changelog

Author: normanrz

webknossos/Changelog.md

@@ -15,6 +15,7 @@ For upgrade instructions, please check the respective _Breaking Changes_ section
 ### Breaking Changes
 
 ### Added
+- Added support for remote datasets and annotations to `webknossos export-as-tiff`. [#1356](https://github.com/scalableminds/webknossos-libs/pull/1356)
 
 ### Changed
 

webknossos/webknossos/cli/export_as_tiff.py

@@ -1,22 +1,29 @@
 """This module takes care of exporting tiff images."""
 
 import logging
+import re
 from argparse import Namespace
 from functools import partial
 from math import ceil
 from multiprocessing import cpu_count
 from pathlib import Path
 from typing import Annotated, Any
+from urllib.parse import urlparse
 
 import numpy as np
 import typer
 from PIL import Image
 from scipy.ndimage import zoom
+from upath import UPath
 
+from ..annotation.annotation import _ANNOTATION_URL_REGEX, Annotation
+from ..client import webknossos_context
+from ..client._resolve_short_link import resolve_short_link
 from ..dataset import Dataset, MagView, View
+from ..dataset.dataset import _DATASET_DEPRECATED_URL_REGEX, _DATASET_URL_REGEX
 from ..dataset.defaults import DEFAULT_CHUNK_SHAPE
 from ..geometry import BoundingBox, Mag, Vec3Int
-from ..utils import get_executor_for_args, wait_and_ensure_success
+from ..utils import get_executor_for_args, is_fs_path, wait_and_ensure_success
 from ._utils import (
     DistributionStrategy,
     Vec2Int,
@@ -173,11 +180,9 @@ def export_tiff_stack(
 def main(
     *,
     source: Annotated[
-        Any,
+        str,
         typer.Argument(
-            help="Path to your raw image data.",
-            show_default=False,
-            parser=parse_path,
+            help="Path or URL to your raw image data.",
         ),
     ],
     target: Annotated[
@@ -258,10 +263,47 @@ def main(
             rich_help_panel="Executor options",
         ),
     ] = None,
+    token: Annotated[
+        str | None,
+        typer.Option(
+            help="Authentication token for WEBKNOSSOS instance "
+            "(https://webknossos.org/auth/token).",
+            rich_help_panel="WEBKNOSSOS context",
+            envvar="WK_TOKEN",
+        ),
+    ] = None,
 ) -> None:
     """Export your WEBKNOSSOS dataset to TIFF image data."""
 
-    mag_view = Dataset.open(source).get_layer(layer_name).get_mag(mag)
+    mag_view = None
+    source_path = UPath(source)
+    if not is_fs_path(source_path):
+        url = resolve_short_link(source)
+        parsed = urlparse(url)
+        domain = f"{parsed.scheme}://{parsed.netloc}"
+
+        with webknossos_context(url=domain, token=token):
+            if re.match(_DATASET_URL_REGEX, url) or re.match(
+                _DATASET_DEPRECATED_URL_REGEX, url
+            ):
+                mag_view = Dataset.open_remote(url).get_layer(layer_name).get_mag(mag)
+            elif re.match(_ANNOTATION_URL_REGEX, url):
+                mag_view = (
+                    Annotation.open_as_remote_dataset(annotation_id_or_url=url)
+                    .get_layer(layer_name)
+                    .get_mag(mag)
+                )
+            else:
+                raise ValueError(
+                    "The provided URL does not lead to a dataset or annotation."
+                )
+    else:
+        mag_view = Dataset.open(source_path).get_layer(layer_name).get_mag(mag)
+
+    if mag_view is None:
+        raise ValueError(
+            f"The provided source does not lead to a dataset or annotation. Got: {source}"
+        )
 
     bbox = BoundingBox.from_ndbbox(mag_view.bounding_box) if bbox is None else bbox
     bbox = bbox.align_with_mag(mag_view.mag)

webknossos/webknossos/dataset/dataset.py

@@ -625,7 +625,7 @@ def _parse_remote(
         sharing_token: str | None = None,
         webknossos_url: str | None = None,
         dataset_id: str | None = None,
-    ) -> tuple[AbstractContextManager, str]:
+    ) -> tuple[AbstractContextManager, str, str | None]:
         """Parses the given arguments to
         * context_manager that should be entered,
         * dataset_id,
@@ -697,10 +697,7 @@ def _parse_remote(
                     + "Please see https://docs.webknossos.org/api/webknossos/client/context.html to adapt the URL and token."
                 )
                 context_manager = webknossos_context(webknossos_url, None)
-        return (
-            context_manager,
-            dataset_id,
-        )
+        return (context_manager, dataset_id, sharing_token)
 
     @classmethod
     def open_remote(
@@ -738,10 +735,7 @@ def open_remote(
 
         from ..client.context import _get_context
 
-        (
-            context_manager,
-            dataset_id,
-        ) = cls._parse_remote(
+        (context_manager, dataset_id, sharing_token) = cls._parse_remote(
             dataset_name_or_url,
             organization_id,
             sharing_token,
@@ -799,10 +793,7 @@ def download(
 
         from ..client._download_dataset import download_dataset
 
-        (
-            context_manager,
-            dataset_id,
-        ) = cls._parse_remote(
+        (context_manager, dataset_id, sharing_token) = cls._parse_remote(
             dataset_name_or_url, organization_id, sharing_token, webknossos_url
         )