Use datastore segment index when loading ad-hoc meshes

frontend/javascripts/admin/rest_api.ts

@@ -855,6 +855,10 @@ export function hasSegmentIndexInDataStore(
   );
 }
 
+export const hasSegmentIndexInDataStoreCached = _.memoize(hasSegmentIndexInDataStore, (...args) =>
+  args.join("::"),
+);
+
 export function getSegmentVolumes(
   requestUrl: string,
   mag: Vector3,
@@ -1927,23 +1931,24 @@ export function computeAdHocMesh(
 }
 
 export function getBucketPositionsForAdHocMesh(
-  tracingStoreUrl: string,
-  tracingId: string,
+  requestUrl: string,
   segmentId: number,
   cubeSize: Vector3,
   mag: Vector3,
   additionalCoordinates: AdditionalCoordinate[] | null | undefined,
+  mappingName: string | null | undefined,
 ): Promise<Vector3[]> {
   return doWithToken(async (token) => {
     const params = new URLSearchParams();
     params.set("token", token);
     const positions = await Request.sendJSONReceiveJSON(
-      `${tracingStoreUrl}/tracings/volume/${tracingId}/segmentIndex/${segmentId}?${params}`,
+      `${requestUrl}/segmentIndex/${segmentId}?${params}`,
       {
         data: {
           cubeSize,
           mag,
           additionalCoordinates,
+          mappingName,
         },
         method: "POST",
       },

frontend/javascripts/viewer/model/sagas/meshes/ad_hoc_mesh_saga.ts

@@ -1,6 +1,7 @@
 import {
   computeAdHocMesh,
   getBucketPositionsForAdHocMesh,
+  hasSegmentIndexInDataStoreCached,
   sendAnalyticsEvent,
 } from "admin/rest_api";
 import ThreeDMap from "libs/ThreeDMap";
@@ -321,6 +322,9 @@ function* loadFullAdHocMesh(
 
   const cubeSize = marchingCubeSizeInTargetMag();
   const tracingStoreHost = yield* select((state) => state.annotation.tracingStore.url);
+  const dataset = yield* select((state) => state.dataset);
+  const datasetId = dataset.id;
+  const dataStoreHost = dataset.dataStore.url;
   const mag = magInfo.getMagByIndexOrThrow(zoomStep);
 
   const volumeTracing = yield* select((state) => getActiveSegmentationTracing(state));
@@ -335,21 +339,39 @@ function* loadFullAdHocMesh(
     useDataStore = false;
   }
 
+  let isStaticSegmentationLayerWithSegmentIndex = false;
+
+  if (volumeTracing == null) {
+    isStaticSegmentationLayerWithSegmentIndex = yield* call(
+      hasSegmentIndexInDataStoreCached,
+      dataset.dataStore.url,
+      dataset.id,
+      layer.name,
+    );
+  }
   // Segment stats can only be used for volume tracings that have a segment index
   // and that don't have editable mappings.
-  const usePositionsFromSegmentIndex =
+  const usePositionsFromSegmentIndexForVolumeTracing =
     volumeTracing?.hasSegmentIndex &&
     !volumeTracing.hasEditableMapping &&
     visibleSegmentationLayer?.tracingId != null;
+
+  const usePositionsFromSegmentIndex =
+    usePositionsFromSegmentIndexForVolumeTracing || isStaticSegmentationLayerWithSegmentIndex;
+
+  const dataStoreUrl = `${dataStoreHost}/data/datasets/${datasetId}/layers/${layer.name}`;
+  const tracingStoreUrl = `${tracingStoreHost}/tracings/volume/${layer.name}`;
+  const requestUrl = useDataStore ? dataStoreUrl : tracingStoreUrl;
+
   let positionsToRequest = usePositionsFromSegmentIndex
     ? yield* getChunkPositionsFromSegmentIndex(
-        tracingStoreHost,
-        layer,
+        requestUrl,
         segmentId,
         cubeSize,
         mag,
         clippedPosition,
         additionalCoordinates,
+        mappingName,
       )
     : [clippedPosition];
 
@@ -390,22 +412,22 @@ function* loadFullAdHocMesh(
 }
 
 function* getChunkPositionsFromSegmentIndex(
-  tracingStoreHost: string,
-  layer: DataLayer,
+  requestUrl: string,
   segmentId: number,
   cubeSize: Vector3,
   mag: Vector3,
   clippedPosition: Vector3,
   additionalCoordinates: AdditionalCoordinate[] | null | undefined,
+  mappingName: string | null | undefined,
 ) {
   const targetMagPositions = yield* call(
     getBucketPositionsForAdHocMesh,
-    tracingStoreHost,
-    layer.name,
+    requestUrl,
     segmentId,
     cubeSize,
     mag,
     additionalCoordinates,
+    mappingName,
   );
   const mag1Positions = targetMagPositions.map((pos) => V3.scale3(pos, mag));
   return sortByDistanceTo(mag1Positions, clippedPosition) as Vector3[];

unreleased_changes/8922.md

@@ -0,0 +1,2 @@
+### Added
+- The segment index file is now used while rendering ad-hoc meshes for static segmentation layers, e.g. when viewing a dataset.

webknossos-datastore/app/com/scalableminds/webknossos/datastore/controllers/BinaryDataController.scala

@@ -224,7 +224,7 @@ class BinaryDataController @Inject()(
           adHocMeshRequest = AdHocMeshRequest(
             Some(dataSource.id),
             segmentationLayer,
-            request.body.cuboid(dataLayer),
+            request.body.cuboid,
             request.body.segmentId,
             request.body.voxelSizeFactorInUnit,
             tokenContextForRequest(request),

webknossos-datastore/app/com/scalableminds/webknossos/datastore/models/DataRequests.scala

@@ -56,7 +56,7 @@ case class WebknossosAdHocMeshRequest(
     additionalCoordinates: Option[Seq[AdditionalCoordinate]] = None,
     findNeighbors: Boolean = true
 ) {
-  def cuboid(dataLayer: DataLayer): Cuboid =
+  def cuboid: Cuboid =
     Cuboid(VoxelPosition(position.x, position.y, position.z, mag), cubeSize.x, cubeSize.y, cubeSize.z)
 }
 

webknossos-datastore/app/com/scalableminds/webknossos/datastore/services/mesh/DSFullMeshService.scala

@@ -13,6 +13,7 @@ import com.scalableminds.webknossos.datastore.services._
 import com.typesafe.scalalogging.LazyLogging
 import com.scalableminds.util.tools.Box.tryo
 import com.scalableminds.webknossos.datastore.services.mapping.MappingService
+import com.scalableminds.webknossos.datastore.services.segmentindex.SegmentIndexFileService
 import play.api.i18n.MessagesProvider
 import play.api.libs.json.{Json, OFormat}
 
@@ -40,6 +41,7 @@ class DSFullMeshService @Inject()(meshFileService: MeshFileService,
                                   val dsRemoteTracingstoreClient: DSRemoteTracingstoreClient,
                                   mappingService: MappingService,
                                   config: DataStoreConfig,
+                                  segmentIndexFileService: SegmentIndexFileService,
                                   adHocMeshServiceHolder: AdHocMeshServiceHolder)
     extends LazyLogging
     with FullMeshHelper
@@ -66,26 +68,85 @@ class DSFullMeshService @Inject()(meshFileService: MeshFileService,
       fullMeshRequest: FullMeshRequest)(implicit ec: ExecutionContext, tc: TokenContext): Fox[Array[Byte]] =
     for {
       mag <- fullMeshRequest.mag.toFox ?~> "mag.neededForAdHoc"
-      seedPosition <- fullMeshRequest.seedPosition.toFox ?~> "seedPosition.neededForAdHoc"
       segmentationLayer <- tryo(dataLayer.asInstanceOf[SegmentationLayer]).toFox ?~> "dataLayer.mustBeSegmentation"
+      hasSegmentIndexFile = segmentationLayer.attachments.flatMap(_.segmentIndex).isDefined
       before = Instant.now
-      verticesForChunks <- getAllAdHocChunks(dataSource,
-                                             segmentationLayer,
-                                             fullMeshRequest,
-                                             VoxelPosition(seedPosition.x, seedPosition.y, seedPosition.z, mag),
-                                             adHocChunkSize)
+      verticesForChunks <- if (hasSegmentIndexFile)
+        getAllAdHocChunksWithSegmentIndex(dataSource, segmentationLayer, fullMeshRequest, mag)
+      else {
+        for {
+          seedPosition <- fullMeshRequest.seedPosition.toFox ?~> "seedPosition.neededForAdHocWithoutSegmentIndex"
+          chunks <- getAllAdHocChunksWithNeighborLogic(
+            dataSource,
+            segmentationLayer,
+            fullMeshRequest,
+            VoxelPosition(seedPosition.x, seedPosition.y, seedPosition.z, mag),
+            adHocChunkSize)
+        } yield chunks
+      }
+
       encoded = verticesForChunks.map(adHocMeshToStl)
       array = combineEncodedChunksToStl(encoded)
       _ = logMeshingDuration(before, "ad-hoc meshing", array.length)
     } yield array
 
-  private def getAllAdHocChunks(
+  private def getAllAdHocChunksWithSegmentIndex(
       dataSource: UsableDataSource,
       segmentationLayer: SegmentationLayer,
       fullMeshRequest: FullMeshRequest,
-      topLeft: VoxelPosition,
-      chunkSize: Vec3Int,
-      visited: collection.mutable.Set[VoxelPosition] = collection.mutable.Set[VoxelPosition]())(
+      mag: Vec3Int,
+  )(implicit ec: ExecutionContext, tc: TokenContext): Fox[List[Array[Float]]] =
+    for {
+      segmentIndexFileKey <- segmentIndexFileService.lookUpSegmentIndexFileKey(dataSource.id, segmentationLayer)
+      segmentIds <- segmentIdsForAgglomerateIdIfNeeded(
+        dataSource.id,
+        segmentationLayer,
+        fullMeshRequest.mappingName,
+        fullMeshRequest.editableMappingTracingId,
+        fullMeshRequest.segmentId,
+        mappingNameForMeshFile = None,
+        omitMissing = false
+      )
+      topLeftsNested: Seq[Array[Vec3Int]] <- Fox.serialCombined(segmentIds)(sId =>
+        segmentIndexFileService.readSegmentIndex(segmentIndexFileKey, sId))
+      topLefts: Array[Vec3Int] = topLeftsNested.toArray.flatten
+      targetMagPositions = segmentIndexFileService.topLeftsToDistinctTargetMagBucketPositions(topLefts, mag)
+      vertexChunksWithNeighbors: List[(Array[Float], List[Int])] <- Fox.serialCombined(targetMagPositions) {
+        targetMagPosition =>
+          val adHocMeshRequest = AdHocMeshRequest(
+            Some(dataSource.id),
+            segmentationLayer,
+            Cuboid(
+              VoxelPosition(
+                targetMagPosition.x * mag.x * DataLayer.bucketLength,
+                targetMagPosition.y * mag.y * DataLayer.bucketLength,
+                targetMagPosition.z * mag.z * DataLayer.bucketLength,
+                mag
+              ),
+              DataLayer.bucketLength + 1,
+              DataLayer.bucketLength + 1,
+              DataLayer.bucketLength + 1
+            ),
+            fullMeshRequest.segmentId,
+            dataSource.scale.factor,
+            tc,
+            fullMeshRequest.mappingName,
+            fullMeshRequest.mappingType,
+            fullMeshRequest.additionalCoordinates,
+            findNeighbors = false,
+          )
+          adHocMeshService.requestAdHocMeshViaActor(adHocMeshRequest)
+      }
+      allVertices = vertexChunksWithNeighbors.map(_._1)
+    } yield allVertices
+
+  private def getAllAdHocChunksWithNeighborLogic(dataSource: UsableDataSource,
+                                                 segmentationLayer: SegmentationLayer,
+                                                 fullMeshRequest: FullMeshRequest,
+                                                 topLeft: VoxelPosition,
+                                                 chunkSize: Vec3Int,
+                                                 visited: collection.mutable.Set[VoxelPosition] =
+                                                   collection.mutable.Set[VoxelPosition]())(
       implicit ec: ExecutionContext,
       tc: TokenContext): Fox[List[Array[Float]]] = {
     val adHocMeshRequest = AdHocMeshRequest(
@@ -105,7 +166,7 @@ class DSFullMeshService @Inject()(meshFileService: MeshFileService,
       nextPositions: List[VoxelPosition] = generateNextTopLeftsFromNeighbors(topLeft, neighbors, chunkSize, visited)
       _ = visited ++= nextPositions
       neighborVerticesNested <- Fox.serialCombined(nextPositions) { position: VoxelPosition =>
-        getAllAdHocChunks(dataSource, segmentationLayer, fullMeshRequest, position, chunkSize, visited)
+        getAllAdHocChunksWithNeighborLogic(dataSource, segmentationLayer, fullMeshRequest, position, chunkSize, visited)
       }
       allVertices: List[Array[Float]] = vertices +: neighborVerticesNested.flatten
     } yield allVertices

webknossos-tracingstore/app/com/scalableminds/webknossos/tracingstore/tracings/editablemapping/EditableMappingService.scala

@@ -401,7 +401,7 @@ class EditableMappingService @Inject()(
     val adHocMeshRequest = AdHocMeshRequest(
       dataSourceId = None,
       dataLayer = editableMappingLayer,
-      cuboid = request.cuboid(editableMappingLayer),
+      cuboid = request.cuboid,
       segmentId = request.segmentId,
       voxelSizeFactor = request.voxelSizeFactorInUnit,
       tokenContext = tc,

webknossos-tracingstore/app/com/scalableminds/webknossos/tracingstore/tracings/volume/VolumeTracingService.scala

@@ -681,7 +681,7 @@ class VolumeTracingService @Inject()(
       adHocMeshRequest = AdHocMeshRequest(
         None,
         volumeLayer,
-        request.cuboid(volumeLayer),
+        request.cuboid,
         request.segmentId,
         request.voxelSizeFactorInUnit,
         tc,