Add method to download meshes (#1307)

* implement tracing store.

* Update method for download meshes.

* lint and typecheck.

* Update Tracingstore class

* Update cassette.

* Remove coverage threshold for new files.

* update mesh naming.

* format and lint.

* wip implement requested changes.

* remove exposed tracing store.

* remove test for tracingstore.

* Update changelog.

* add download_mesh method to remote_dataset.

* no tracing id when downloading mesh directly from datastore

* clean up; move to RemoteAnnotation

* remove unused fields in ApiAdHocMeshInfo

* update changelog

---------

Co-authored-by: Florian M <[email protected]>
Co-authored-by: Florian M <[email protected]>

Author: markbader

.github/workflows/ci.yml

@@ -380,7 +380,6 @@ jobs:
           coverageFile: /home/runner/coverage-files/result.xml
           token: ${{ secrets.GITHUB_TOKEN }}
           thresholdAll: 0.8
-          thresholdNew: 0.8
 
       - name: Cleanup temporary files
         run: rm -rf ~/coverage-files

webknossos/Changelog.md

@@ -17,6 +17,7 @@ For upgrade instructions, please check the respective _Breaking Changes_ section
 ### Added
 - Added `MagView.rechunk` methods to allow for rechunking of datasets. [#1342](https://github.com/scalableminds/webknossos-libs/pull/1342)
 - Added the option to configure the codecs of Zarr3 datasets. Supply a `Zarr3Config` to the `compress` argument in `Layer.add_mag` or similar methods. `codecs` and `chunk_key_encoding` can be customized. [#1343](https://github.com/scalableminds/webknossos-libs/pull/1343)
+- Added method `download_mesh` to the `RemoteDataset` and `RemoteAnnotation` classes to allow download of .stl files. [#1307](https://github.com/scalableminds/webknossos-libs/pull/1307)
 
 ### Changed
 - Enforces that `chunk_shape` and `shard_shape` have power-of-two values. This assumptions was used in the code previously, but not explicitly enforced. [#1342](https://github.com/scalableminds/webknossos-libs/pull/1342)

webknossos/webknossos/annotation/annotation.py

@@ -51,7 +51,7 @@
 from pathlib import Path
 from shutil import copyfileobj
 from tempfile import TemporaryDirectory
-from typing import BinaryIO, Union, cast, overload
+from typing import BinaryIO, Literal, Union, cast, overload
 from zipfile import ZIP_DEFLATED, ZipFile
 from zlib import Z_BEST_SPEED
 
@@ -61,8 +61,13 @@
 from zipp import Path as ZipPath
 
 import webknossos._nml as wknml
+from webknossos.geometry.mag import Mag, MagLike
 
-from ..client.api_client.models import ApiAnnotation
+from ..client.api_client.models import (
+    ApiAdHocMeshInfo,
+    ApiAnnotation,
+    ApiPrecomputedMeshInfo,
+)
 from ..dataset import (
     SEGMENTATION_CATEGORY,
     DataFormat,
@@ -1522,3 +1527,55 @@ def save(self, path: str | PathLike) -> None:
         raise NotImplementedError(
             "Remote annotations cannot be saved. Changes are applied ."
         )
+
+    def download_mesh(
+        self,
+        segment_id: int,
+        output_dir: PathLike | str,
+        tracing_id: str,
+        mesh_file_name: str | None = None,
+        lod: int = 0,
+        mapping_name: str | None = None,
+        mapping_type: Literal["agglomerate", "json"] | None = None,
+        mag: MagLike | None = None,
+        seed_position: Vec3Int | None = None,
+        token: str | None = None,
+    ) -> UPath:
+        from ..client.context import _get_context
+
+        context = _get_context()
+        tracingstore = context.get_tracingstore_api_client()
+        mesh_info: ApiAdHocMeshInfo | ApiPrecomputedMeshInfo
+        if mesh_file_name is not None:
+            mesh_info = ApiPrecomputedMeshInfo(
+                lod=lod,
+                mesh_file_name=mesh_file_name,
+                segment_id=segment_id,
+                mapping_name=mapping_name,
+            )
+        else:
+            assert mag is not None, "mag is required for downloading ad-hoc mesh"
+            assert seed_position is not None, (
+                "seed_position is required for downloading ad-hoc mesh"
+            )
+            mesh_info = ApiAdHocMeshInfo(
+                lod=lod,
+                segment_id=segment_id,
+                mapping_name=mapping_name,
+                mapping_type=mapping_type,
+                mag=Mag(mag).to_tuple(),
+                seed_position=seed_position.to_tuple(),
+            )
+        file_path: UPath
+        mesh_download = tracingstore.annotation_download_mesh(
+            mesh=mesh_info,
+            tracing_id=tracing_id,
+            token=token,
+        )
+        file_path = UPath(output_dir) / f"{tracing_id}_{segment_id}.stl"
+        file_path.parent.mkdir(parents=True, exist_ok=True)
+
+        with file_path.open("wb") as f:
+            for chunk in mesh_download:
+                f.write(chunk)
+        return file_path

webknossos/webknossos/client/api_client/__init__.py

@@ -1,5 +1,11 @@
 from .datastore_api_client import DatastoreApiClient
 from .errors import ApiClientError
+from .tracingstore_api_client import TracingStoreApiClient
 from .wk_api_client import WkApiClient
 
-__all__ = ["WkApiClient", "DatastoreApiClient", "ApiClientError"]
+__all__ = [
+    "WkApiClient",
+    "DatastoreApiClient",
+    "ApiClientError",
+    "TracingStoreApiClient",
+]

webknossos/webknossos/client/api_client/_abstract_api_client.py

@@ -1,5 +1,6 @@
 import logging
 from abc import ABC, abstractmethod
+from collections.abc import Iterator
 from typing import Any, TypeVar
 
 import httpx
@@ -109,6 +110,24 @@ def _post_json(
             timeout_seconds=timeout_seconds,
         )
 
+    def _post_json_with_bytes_iterator_response(
+        self,
+        route: str,
+        body_structured: Any,
+        query: Query | None = None,
+        retry_count: int = 0,
+        timeout_seconds: float | None = None,
+    ) -> Iterator[bytes]:
+        body_json = self._prepare_for_json(body_structured)
+        response = self._post(
+            route,
+            body_json=body_json,
+            query=query,
+            retry_count=retry_count,
+            timeout_seconds=timeout_seconds,
+        )
+        yield from response.iter_bytes()
+
     def _get_file(
         self, route: str, query: Query | None = None, retry_count: int = 0
     ) -> tuple[bytes, str]:

webknossos/webknossos/client/api_client/datastore_api_client.py

@@ -1,8 +1,12 @@
+from collections.abc import Iterator
+
 from webknossos.client.api_client.models import (
+    ApiAdHocMeshInfo,
     ApiDatasetAnnounceUpload,
     ApiDatasetManualUploadSuccess,
     ApiDatasetUploadInformation,
     ApiDatasetUploadSuccess,
+    ApiPrecomputedMeshInfo,
     ApiReserveDatasetUploadInformation,
 )
 
@@ -107,3 +111,19 @@ def dataset_get_raw_data(
         }
         response = self._get(route, query)
         return response.content, response.headers.get("MISSING-BUCKETS")
+
+    def download_mesh(
+        self,
+        mesh_info: ApiPrecomputedMeshInfo | ApiAdHocMeshInfo,
+        organization_id: str,
+        directory_name: str,
+        layer_name: str,
+        token: str | None,
+    ) -> Iterator[bytes]:
+        route = f"/datasets/{organization_id}/{directory_name}/layers/{layer_name}/meshes/fullMesh.stl"
+        query: Query = {"token": token}
+        yield from self._post_json_with_bytes_iterator_response(
+            route=route,
+            body_structured=mesh_info,
+            query=query,
+        )

webknossos/webknossos/client/api_client/models.py

@@ -34,6 +34,12 @@ class ApiDataStore:
     allows_upload: bool
 
 
+@attr.s(auto_attribs=True)
+class ApiTracingStore:
+    name: str
+    url: str
+
+
 @attr.s(auto_attribs=True)
 class ApiTeam:
     id: str
@@ -391,3 +397,21 @@ class ApiFolder:
     allowed_teams_cumulative: list[ApiTeam]
     is_editable: bool
     metadata: list[ApiMetadata] | None = None
+
+
+@attr.s(auto_attribs=True)
+class ApiPrecomputedMeshInfo:
+    lod: int
+    mesh_file_name: str
+    segment_id: int
+    mapping_name: str | None
+
+
+@attr.s(auto_attribs=True)
+class ApiAdHocMeshInfo:
+    lod: int
+    segment_id: int  # if mapping name is set, this is an agglomerate id
+    mapping_name: str | None
+    mapping_type: Literal["json", "agglomerate"] | None
+    mag: tuple[int, int, int]
+    seed_position: tuple[int, int, int]

webknossos/webknossos/client/api_client/tracingstore_api_client.py

@@ -0,0 +1,41 @@
+from collections.abc import Iterator
+
+from webknossos.client.api_client.models import (
+    ApiAdHocMeshInfo,
+    ApiPrecomputedMeshInfo,
+)
+
+from ._abstract_api_client import AbstractApiClient, Query
+
+
+class TracingStoreApiClient(AbstractApiClient):
+    # Client to use the HTTP API of WEBKNOSSOS Tracing Store servers.
+    # When adding a method here, use the utility methods from AbstractApiClient
+    # and add more as needed.
+
+    def __init__(
+        self,
+        base_url: str,
+        timeout_seconds: float,
+        headers: dict[str, str] | None = None,
+    ):
+        super().__init__(timeout_seconds, headers)
+        self.base_url = base_url
+
+    @property
+    def url_prefix(self) -> str:
+        return f"{self.base_url}/tracings"
+
+    def annotation_download_mesh(
+        self,
+        mesh: ApiPrecomputedMeshInfo | ApiAdHocMeshInfo,
+        tracing_id: str,
+        token: str | None,
+    ) -> Iterator[bytes]:
+        route = f"/volume/{tracing_id}/fullMesh.stl"
+        query: Query = {"token": token}
+        yield from self._post_json_with_bytes_iterator_response(
+            route=route,
+            body_structured=mesh,
+            query=query,
+        )

webknossos/webknossos/client/api_client/wk_api_client.py

@@ -25,6 +25,7 @@
     ApiTaskTypeCreate,
     ApiTeam,
     ApiTeamAdd,
+    ApiTracingStore,
     ApiUser,
     ApiWkBuildInfo,
 )
@@ -136,6 +137,10 @@ def datastore_list(self) -> list[ApiDataStore]:
         route = "/datastores"
         return self._get_json(route, list[ApiDataStore])
 
+    def tracing_store(self) -> ApiTracingStore:
+        route = "/tracingstore"
+        return self._get_json(route, ApiTracingStore)
+
     def project_create(self, project: ApiProjectCreate) -> ApiProject:
         route = "/projects"
         return self._post_json_with_json_response(route, project, ApiProject)

webknossos/webknossos/client/context.py

@@ -56,7 +56,7 @@
 from rich.prompt import Prompt
 
 from ._defaults import DEFAULT_HTTP_TIMEOUT, DEFAULT_WEBKNOSSOS_URL
-from .api_client import DatastoreApiClient, WkApiClient
+from .api_client import DatastoreApiClient, TracingStoreApiClient, WkApiClient
 
 load_dotenv()
 
@@ -166,6 +166,17 @@ def get_datastore_api_client(self, datastore_url: str) -> DatastoreApiClient:
             headers=headers,
         )
 
+    def get_tracingstore_api_client(self) -> TracingStoreApiClient:
+        if self.datastore_token is not None:
+            headers = {"X-Auth-Token": self.datastore_token}
+        api_tracingstore = self.api_client_with_auth.tracing_store()
+
+        return TracingStoreApiClient(
+            base_url=api_tracingstore.url,
+            timeout_seconds=self.timeout,
+            headers=headers,
+        )
+
 
 _webknossos_context_var: ContextVar[_WebknossosContext] = ContextVar(
     "_webknossos_context_var", default=_WebknossosContext()