feat(segmentation_user_layer) implemented getObjectPosition for MeshL…

…ayer, finds the closest loaded associated mesh vertex to the global position
seung-lab · Feb 25, 2024 · c067351 · c067351
1 parent 09659df
commit c067351
Show file tree

Hide file tree

Showing 2 changed files with 106 additions and 17 deletions.
diff --git a/src/layer/segmentation/index.ts b/src/layer/segmentation/index.ts
@@ -41,7 +41,6 @@ import {
   MultiscaleMeshLayer,
   MultiscaleMeshSource,
 } from "#src/mesh/frontend.js";
-import type { RenderLayerTransform } from "#src/render_coordinate_transform.js";
 import {
   RenderScaleHistogram,
   trackableRenderScaleTarget,
@@ -130,11 +129,14 @@ import { Uint64 } from "#src/util/uint64.js";
 import { makeWatchableShaderError } from "#src/webgl/dynamic_shader.js";
 import type { DependentViewContext } from "#src/widget/dependent_view_widget.js";
 import { registerLayerShaderControlsTool } from "#src/widget/shader_controls.js";
+import { gatherUpdate } from "#src/util/array";
+
+export const SKELETON_RENDERING_SHADER_CONTROL_TOOL_ID =
+  "skeletonShaderControl";
 
 export class SegmentationUserLayerGroupState
   extends RefCounted
-  implements SegmentationGroupState
-{
+  implements SegmentationGroupState {
   specificationChanged = new Signal();
   constructor(public layer: SegmentationUserLayer) {
     super();
@@ -281,8 +283,7 @@ export class SegmentationUserLayerGroupState
 
 export class SegmentationUserLayerColorGroupState
   extends RefCounted
-  implements SegmentationColorGroupState
-{
+  implements SegmentationColorGroupState {
   specificationChanged = new Signal();
   constructor(public layer: SegmentationUserLayer) {
     super();
@@ -356,10 +357,10 @@ export class SegmentationUserLayerColorGroupState
 }
 
 class LinkedSegmentationGroupState<
-    State extends
-      | SegmentationUserLayerGroupState
-      | SegmentationUserLayerColorGroupState,
-  >
+  State extends
+  | SegmentationUserLayerGroupState
+  | SegmentationUserLayerColorGroupState,
+>
   extends RefCounted
   implements WatchableValueInterface<State>
 {
@@ -786,7 +787,7 @@ export class SegmentationUserLayer extends Base {
             "Not supported on non-root linked segmentation layers",
           );
         } else {
-          loadedSubsource.activate(() => {});
+          loadedSubsource.activate(() => { });
           updatedSegmentPropertyMaps.push(segmentPropertyMap);
         }
       } else if (segmentationGraph !== undefined) {
@@ -913,7 +914,7 @@ export class SegmentationUserLayer extends Base {
     if (
       layerSpec[json_keys.EQUIVALENCES_JSON_KEY] !== undefined &&
       explicitSpecs.find((spec) => spec.url === localEquivalencesUrl) ===
-        undefined
+      undefined
     ) {
       specs.push({
         url: localEquivalencesUrl,
@@ -1229,13 +1230,31 @@ export class SegmentationUserLayer extends Base {
 
   moveToSegment(id: Uint64) {
     for (const layer of this.renderLayers) {
-      if (!(layer instanceof MultiscaleMeshLayer)) continue;
-      const layerPosition = layer.getObjectPosition(id);
-      if (layerPosition === undefined) continue;
-      this.setLayerPosition(
-        layer.displayState.transform.value as RenderLayerTransform,
-        layerPosition,
+      if (
+        !(layer instanceof MultiscaleMeshLayer || layer instanceof MeshLayer)
+      ) {
+        continue;
+      }
+      const transform = layer.displayState.transform.value;
+      if (transform.error !== undefined) return undefined;
+      const { rank, globalToRenderLayerDimensions } = transform;
+      const { globalPosition } = this.manager.root;
+      const globalLayerPosition = new Float32Array(rank);
+      const renderToGlobalLayerDimensions = [];
+      for (let i = 0; i < rank; i++) {
+        renderToGlobalLayerDimensions[globalToRenderLayerDimensions[i]] = i;
+      }
+      gatherUpdate(
+        globalLayerPosition,
+        globalPosition.value,
+        renderToGlobalLayerDimensions,
       );
+      const layerPosition =
+        layer instanceof MeshLayer
+          ? layer.getObjectPosition(id, globalLayerPosition)
+          : layer.getObjectPosition(id);
+      if (layerPosition === undefined) continue;
+      this.setLayerPosition(transform, layerPosition);
       return;
     }
     StatusMessage.showTemporaryMessage(

diff --git a/src/mesh/frontend.ts b/src/mesh/frontend.ts
@@ -582,6 +582,76 @@ export class MeshLayer extends PerspectiveViewRenderLayer<ThreeDimensionalRender
     );
     return ready;
   }
+
+  getObjectPosition(
+    id: Uint64,
+    nearestTo: Float32Array,
+  ): Float32Array | undefined {
+    const transform = this.displayState.transform.value;
+    if (transform.error !== undefined) return undefined;
+    const chunk = this.source.chunks.get(getObjectKey(id));
+    if (chunk === undefined) return undefined;
+    const { rank } = transform;
+    const inverseModelToRenderLayerTransform = new Float32Array(
+      transform.modelToRenderLayerTransform.length,
+    );
+    matrix.inverse(
+      inverseModelToRenderLayerTransform,
+      rank + 1,
+      transform.modelToRenderLayerTransform,
+      rank + 1,
+      rank + 1,
+    );
+    const nearestPositionInModal = new Float32Array(rank);
+    matrix.transformPoint(
+      nearestPositionInModal,
+      inverseModelToRenderLayerTransform,
+      rank + 1,
+      nearestTo,
+      rank,
+    );
+    const { fragmentIds } = chunk;
+    let closestVertex = new Float32Array(rank);
+    let closestDistanceSq = Number.POSITIVE_INFINITY;
+    for (let fragmentId of fragmentIds) {
+      const fragmentChunk = this.source.fragmentSource.chunks.get(fragmentId);
+      if (fragmentChunk === undefined) continue;
+      const { state, meshData } = fragmentChunk;
+      if (
+        state !== ChunkState.SYSTEM_MEMORY &&
+        state !== ChunkState.GPU_MEMORY
+      ) {
+        continue;
+      }
+      const { vertexPositions } = meshData;
+      if (vertexPositions.length < rank) continue;
+      for (let i = 0; i < vertexPositions.length; i += rank) {
+        let distanceSq = 0;
+        for (let j = 0; j < rank; j++) {
+          distanceSq += Math.pow(
+            vertexPositions[i + j] - nearestPositionInModal[j],
+            2,
+          );
+        }
+        if (distanceSq < closestDistanceSq) {
+          closestDistanceSq = distanceSq;
+          for (let j = 0; j < rank; j++) {
+            closestVertex[j] = vertexPositions[i + j];
+          }
+        }
+      }
+    }
+    if (closestDistanceSq === Number.POSITIVE_INFINITY) return undefined;
+    const layerCenter = new Float32Array(rank);
+    matrix.transformPoint(
+      layerCenter,
+      transform.modelToRenderLayerTransform,
+      rank + 1,
+      closestVertex,
+      rank,
+    );
+    return layerCenter;
+  }
 }
 
 export class ManifestChunk extends Chunk {