refactor dataset, remote_dataset, and layer

webknossos/examples/dataset_usage.py

@@ -9,7 +9,6 @@ def main() -> None:
     #####################
     # Opening a dataset #
     #####################
-
     dataset = wk.Dataset.open("testdata/simple_wkw_dataset")
     # Assuming that the dataset has a layer "color"
     # and the layer has the magnification 1

webknossos/webknossos/administration/task.py

@@ -5,6 +5,7 @@
 
 from webknossos.geometry.vec3_int import Vec3IntLike
 
+from .. import RemoteDataset
 from ..annotation import Annotation, AnnotationInfo
 from ..client.api_client.models import (
     ApiBoundingBox,
@@ -17,7 +18,6 @@
     ApiTaskTypeCreate,
 )
 from ..client.context import _get_api_client, _get_context
-from ..dataset.dataset import RemoteDataset
 from ..geometry import BoundingBox, Vec3Int
 from ..utils import warn_deprecated
 from .project import Project

webknossos/webknossos/annotation/annotation.py

@@ -63,6 +63,7 @@
 import webknossos._nml as wknml
 from webknossos.geometry.mag import Mag, MagLike
 
+from .. import RemoteDataset
 from ..client.api_client.models import (
     ApiAdHocMeshInfo,
     ApiAnnotation,
@@ -71,7 +72,6 @@
 from ..dataset import (
     Dataset,
     Layer,
-    RemoteDataset,
     SegmentationLayer,
 )
 from ..dataset.defaults import PROPERTIES_FILE_NAME
@@ -84,7 +84,6 @@
 from ..dataset_properties.structuring import get_dataset_converter
 from ..geometry import NDBoundingBox, Vec3Int
 from ..skeleton import Skeleton
-from ..ssl_context import SSL_CONTEXT
 from ..utils import get_executor_for_args, time_since_epoch_in_ms
 from ._nml_conversion import annotation_to_nml, nml_to_skeleton
 
@@ -532,7 +531,7 @@ def open_as_remote_dataset(
         cls,
         annotation_id_or_url: str,
         webknossos_url: str | None = None,
-    ) -> Dataset:
+    ) -> RemoteDataset:
         """Opens an annotation directly as a remote dataset from WEBKNOSSOS.
 
         This is a convenience method that combines downloading an annotation and converting it
@@ -926,7 +925,7 @@ def _write_to_zip(self, zipfile: ZipFile) -> None:
                 layer_content,
             )
 
-    def get_remote_annotation_dataset(self) -> Dataset:
+    def get_remote_annotation_dataset(self) -> RemoteDataset:
         """Returns a streamed dataset of the annotation from WEBKNOSSOS.
 
         Creates a remote dataset that includes fallback layers and potentially any active agglomerate mappings.
@@ -962,7 +961,8 @@ def get_remote_annotation_dataset(self) -> Dataset:
         context = _get_context()
         token: str | None
         if self.organization_id is None:
-            token = context.required_token
+            # ask for token, if necessary
+            _ = context.required_token
             organization_id = context.organization_id
         else:
             token = context.token
@@ -976,24 +976,15 @@ def get_remote_annotation_dataset(self) -> Dataset:
                         + f"The annotation uses {organization_id}, the context {context.organization_id}.",
                         UserWarning,
                     )
-        if self.dataset_id is None:
-            dataset_id = context.api_client.dataset_id_from_name(
-                directory_name=self.dataset_name, organization_id=organization_id
-            )
-        else:
-            dataset_id = self.dataset_id
-
-        dataset_info = context.api_client.dataset_info(dataset_id=dataset_id)
 
-        datastore_url = dataset_info.data_store.url
-        url_prefix = context.get_datastore_api_client(datastore_url).url_prefix
-
-        zarr_path = UPath(
-            f"{url_prefix}/annotations/zarr/{self.annotation_id}/",
-            headers={} if token is None else {"X-Auth-Token": token},
-            ssl=SSL_CONTEXT,
+        return RemoteDataset.open(
+            dataset_name_or_url=self.dataset_name,
+            organization_id=organization_id,
+            dataset_id=self.dataset_id,
+            annotation_id=self.annotation_id,
+            use_zarr_streaming=True,
+            read_only=True,
         )
-        return Dataset.open(zarr_path)
 
     def get_remote_base_dataset(
         self,

webknossos/webknossos/cli/convert_raw.py

@@ -10,6 +10,7 @@
 import numpy as np
 import typer
 
+from .. import Layer
 from ..dataset import Dataset, MagView, SamplingModes
 from ..dataset.defaults import (
     DEFAULT_CHUNK_SHAPE,
@@ -81,7 +82,7 @@ def convert_raw(
     flip_axes: int | tuple[int, ...] | None = None,
     compress: bool = True,
     executor_args: Namespace | None = None,
-) -> MagView:
+) -> tuple[MagView, Layer]:
     """Performs the conversion step from RAW file to WEBKNOSSOS"""
     time_start(f"Conversion of {source_raw_path}")
 
@@ -122,7 +123,7 @@ def convert_raw(
         )
 
     time_stop(f"Conversion of {source_raw_path}")
-    return wk_mag
+    return wk_mag, wk_layer
 
 
 def main(
@@ -297,7 +298,7 @@ def main(
     )
     voxel_size_with_unit = VoxelSize(voxel_size, unit)
 
-    mag_view = convert_raw(
+    mag_view, layer = convert_raw(
         source,
         target,
         layer_name,
@@ -314,7 +315,7 @@ def main(
     )
 
     with get_executor_for_args(executor_args) as executor:
-        mag_view.layer.downsample(
+        layer.downsample(
             from_mag=mag_view.mag,
             coarsest_mag=max_mag,
             interpolation_mode=interpolation_mode,

webknossos/webknossos/cli/convert_zarr.py

@@ -12,6 +12,7 @@
 import tensorstore
 import typer
 
+from .. import Layer
 from ..dataset import Dataset, MagView, SegmentationLayer
 from ..dataset.defaults import (
     DEFAULT_CHUNK_SHAPE,
@@ -80,7 +81,7 @@ def convert_zarr(
     flip_axes: int | tuple[int, ...] | None = None,
     compress: bool = True,
     executor_args: Namespace | None = None,
-) -> MagView:
+) -> tuple[MagView, Layer]:
     """Performs the conversation of a Zarr dataset to a WEBKNOSSOS dataset."""
     ref_time = time.time()
 
@@ -129,7 +130,7 @@ def convert_zarr(
     logger.debug(
         "Conversion of %s took %.8fs", source_zarr_path, time.time() - ref_time
     )
-    return wk_mag
+    return wk_mag, wk_layer
 
 
 def main(
@@ -283,7 +284,7 @@ def main(
     )
     voxel_size_with_unit = VoxelSize(voxel_size, unit)
 
-    mag_view = convert_zarr(
+    mag_view, layer = convert_zarr(
         source,
         target,
         layer_name=layer_name,
@@ -298,7 +299,7 @@ def main(
     )
 
     with get_executor_for_args(executor_args) as executor:
-        mag_view.layer.downsample(
+        layer.downsample(
             from_mag=mag_view.mag,
             coarsest_mag=max_mag,
             interpolation_mode=interpolation_mode,

webknossos/webknossos/cli/download.py

@@ -9,7 +9,8 @@
 from ..annotation.annotation import _ANNOTATION_URL_REGEX, Annotation
 from ..client import webknossos_context
 from ..client._resolve_short_link import resolve_short_link
-from ..dataset.dataset import _DATASET_DEPRECATED_URL_REGEX, _DATASET_URL_REGEX, Dataset
+from ..dataset import Dataset
+from ..dataset.abstract_dataset import _DATASET_DEPRECATED_URL_REGEX, _DATASET_URL_REGEX
 from ..geometry import BoundingBox, Mag
 from ._utils import parse_bbox, parse_mag, parse_path
 

webknossos/webknossos/cli/export_as_tiff.py

@@ -20,7 +20,7 @@
 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.abstract_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, is_fs_path, wait_and_ensure_success

webknossos/webknossos/client/_upload_dataset.py

@@ -9,7 +9,8 @@
 
 import httpx
 
-from ..dataset import Dataset, LayerToLink
+from .. import LayerToLink
+from ..dataset import Dataset
 from ..datastore import Datastore
 from ..utils import get_rich_progress
 from ._resumable import Resumable

webknossos/webknossos/dataset/__init__.py

@@ -18,15 +18,17 @@
     MeshAttachment,
     SegmentIndexAttachment,
 )
-from .dataset import Dataset, LayerToLink, RemoteDataset
+from .dataset import Dataset
 from .defaults import (
     DEFAULT_CHUNK_SHAPE,
     DEFAULT_CHUNKS_PER_SHARD_FROM_IMAGES,
     DEFAULT_DATA_FORMAT,
     DEFAULT_SHARD_SHAPE,
 )
 from .layer import Layer, SegmentationLayer
+from .layer_to_link import LayerToLink
 from .mag_view import MagView
+from .remote_dataset import RemoteDataset
 from .remote_folder import RemoteFolder
 from .sampling_modes import SamplingModes
 from .view import View