segment color shader

Author: chrisj

src/layer/segmentation/index.ts

@@ -129,9 +129,16 @@ import {
   verifyString,
 } from "#src/util/json.js";
 import { Signal } from "#src/util/signal.js";
-import { makeWatchableShaderError } from "#src/webgl/dynamic_shader.js";
+import {
+  makeTrackableFragmentMain,
+  makeWatchableShaderError,
+  parameterizedEmitterDependentShaderGetter,
+} 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 { ShaderControlState } from "#src/webgl/shader_ui_controls.js";
+import { initializeWebGL } from "#src/webgl/context.js";
+import { GLBuffer } from "#src/webgl/buffer.js";
 
 const MAX_LAYER_BAR_UI_INDICATOR_COLORS = 6;
 
@@ -410,12 +417,22 @@ class LinkedSegmentationGroupState<
   }
 }
 
+const DEFAULT_FRAGMENT_MAIN = `
+vec4 segmentColor(vec4 color) {
+  return color;
+}
+`;
+
 class SegmentationUserLayerDisplayState implements SegmentationDisplayState {
+  private getSegmentColorShader;
+
   constructor(public layer: SegmentationUserLayer) {
     // Even though `SegmentationUserLayer` assigns this to its `displayState` property, redundantly
     // assign it here first in order to allow it to be accessed by `segmentationGroupState`.
     layer.displayState = this;
 
+    this.getSegmentColorShader = this.makeSegmentColorShaderGetter();
+
     this.linkedSegmentationGroup = layer.registerDisposer(
       new LinkedLayerGroup(
         layer.manager.rootLayers,
@@ -536,6 +553,8 @@ class SegmentationUserLayerDisplayState implements SegmentationDisplayState {
   ignoreNullVisibleSet = new TrackableBoolean(true, true);
   skeletonRenderingOptions = new SkeletonRenderingOptions();
   shaderError = makeWatchableShaderError();
+  fragmentMain = makeTrackableFragmentMain(DEFAULT_FRAGMENT_MAIN);
+  shaderControlState = new ShaderControlState(this.fragmentMain);
   renderScaleHistogram = new RenderScaleHistogram();
   renderScaleTarget = trackableRenderScaleTarget(1);
   selectSegment: (id: bigint, pin: boolean | "toggle") => void;
@@ -551,6 +570,67 @@ class SegmentationUserLayerDisplayState implements SegmentationDisplayState {
     this.layer.moveToSegment(id);
   };
 
+  makeSegmentColorShaderGetter = () => {
+    const gl = initializeWebGL(new OffscreenCanvas(1, 1));
+    const parameters = this.shaderControlState.builderState;
+    return parameterizedEmitterDependentShaderGetter(this.layer, gl, {
+      memoizeKey: `segmentColorShaderTODO`,
+      parameters,
+      encodeParameters: (p) => {
+        return `${p.parseResult.code}`;
+      },
+      shaderError: this.layer.displayState.shaderError, // TODO can I reuse this?
+      defineShader: (builder, shaderBuilderState) => {
+        builder.addAttribute("highp vec4", "aVertexPosition");
+        builder.addUniform("highp vec4", "uColor");
+        builder.addUniform("highp uvec2", "uID");
+        builder.addVarying("highp vec4", "vColor");
+        builder.addVertexCode(
+          shaderBuilderState.parseResult.code
+        );
+        builder.addVertexMain(`
+gl_Position = aVertexPosition;
+vColor = segmentColor(uColor);
+`);
+        builder.addOutputBuffer("vec4", "out_fragColor", 0);
+        builder.setFragmentMain("out_fragColor = vColor;");
+      },
+    });
+  };
+
+  context = () => {}; // TEMP how to get rid of this?
+
+  getShaderSegmentColor = (id: bigint, color: Float32Array) => {
+    id; // TODO
+    const { shader, parameters } = this.getSegmentColorShader(this.context);
+    if (shader === null) return color;
+    parameters;
+    shader.bind();
+    const { gl } = shader;
+    const positionBuffer = GLBuffer.fromData(
+      gl,
+      new Float32Array([1, 1, -1, 1, 1, -1, -1, -1]),
+    );
+    positionBuffer.bindToVertexAttrib(shader.attribute("aVertexPosition"), 2);
+    gl.uniform4fv(shader.uniform("uColor"), color);
+    gl.drawArrays(gl.TRIANGLE_STRIP, 0, 4);
+    const data = new Uint8Array(4);
+    // TODO can I read straight to float?
+    gl.readPixels(
+      0,
+      0,
+      1,
+      1,
+      WebGL2RenderingContext.RGBA,
+      WebGL2RenderingContext.UNSIGNED_BYTE,
+      data,
+    );
+    for (let i = 0; i < data.length; i++) {
+      color[i] = data[i] / 255.0;
+    }
+    return color;
+  };
+
   linkedSegmentationGroup: LinkedLayerGroup;
   linkedSegmentationColorGroup: LinkedLayerGroup;
   originalSegmentationGroupState: SegmentationUserLayerGroupState;
@@ -689,6 +769,9 @@ export class SegmentationUserLayer extends Base {
     this.displayState.linkedSegmentationGroup.changed.add(() =>
       this.updateDataSubsourceActivations(),
     );
+    this.displayState.fragmentMain.changed.add(
+      this.specificationChanged.dispatch,
+    );
     this.tabs.add("rendering", {
       label: "Render",
       order: -100,
@@ -996,6 +1079,9 @@ export class SegmentationUserLayer extends Base {
     this.displayState.ignoreNullVisibleSet.restoreState(
       specification[json_keys.IGNORE_NULL_VISIBLE_SET_JSON_KEY],
     );
+    this.displayState.fragmentMain.restoreState(
+      specification[json_keys.SHADER_JSON_KEY],
+    );
 
     const { skeletonRenderingOptions } = this.displayState;
     skeletonRenderingOptions.restoreState(
@@ -1066,6 +1152,7 @@ export class SegmentationUserLayer extends Base {
       this.displayState.renderScaleTarget.toJSON();
     x[json_keys.CROSS_SECTION_RENDER_SCALE_JSON_KEY] =
       this.sliceViewRenderScaleTarget.toJSON();
+    x[json_keys.SHADER_JSON_KEY] = this.displayState.fragmentMain.toJSON();
 
     const { linkedSegmentationGroup, linkedSegmentationColorGroup } =
       this.displayState;
@@ -1365,6 +1452,7 @@ export class SegmentationUserLayer extends Base {
 
     const visibleSegments = [...visibleSegmentsSet];
     const colors = visibleSegments.map((id) => {
+      // here we can do a batch get of colors using the segment color shader instead of one at a time
       const color = getCssColor(getBaseObjectColor(displayState, id));
       return { color, id };
     });
@@ -1401,6 +1489,10 @@ registerLayerTypeDetector((subsource) => {
   return undefined;
 });
 
+registerLayerShaderControlsTool(SegmentationUserLayer, (layer) => ({
+  shaderControlState: layer.displayState.shaderControlState,
+}));
+
 registerLayerShaderControlsTool(
   SegmentationUserLayer,
   (layer) => ({

src/layer/segmentation/json_keys.ts

@@ -14,6 +14,7 @@ export const COLOR_SEED_JSON_KEY = "colorSeed";
 export const SEGMENT_STATED_COLORS_JSON_KEY = "segmentColors";
 export const MESH_RENDER_SCALE_JSON_KEY = "meshRenderScale";
 export const CROSS_SECTION_RENDER_SCALE_JSON_KEY = "crossSectionRenderScale";
+export const SHADER_JSON_KEY = "shader";
 export const SKELETON_RENDERING_JSON_KEY = "skeletonRendering";
 export const SKELETON_SHADER_JSON_KEY = "skeletonShader";
 export const SKELETON_CODE_VISIBLE_KEY = "codeVisible";

src/layer/segmentation/style.css

@@ -22,6 +22,7 @@
 
 .neuroglancer-segmentation-rendering-tab .neuroglancer-shader-code-widget {
   height: 6em;
+  resize: vertical;
 }
 
 .neuroglancer-segmentation-dropdown-skeleton-shader-header {

src/mesh/frontend.ts

@@ -43,6 +43,9 @@ import type {
 import { PerspectiveViewRenderLayer } from "#src/perspective_view/render_layer.js";
 import type { ThreeDimensionalRenderLayerAttachmentState } from "#src/renderlayer.js";
 import { update3dRenderLayerAttachment } from "#src/renderlayer.js";
+import {
+  SegmentColorShaderManager,
+} from "#src/segment_color.js";
 import {
   forEachVisibleSegment,
   getObjectKey,
@@ -56,7 +59,7 @@ import {
 import type { WatchableValueInterface } from "#src/trackable_value.js";
 import { makeCachedDerivedWatchableValue } from "#src/trackable_value.js";
 import type { Borrowed, RefCounted } from "#src/util/disposable.js";
-import type { vec4 } from "#src/util/geom.js";
+import { vec4 } from "#src/util/geom.js";
 import {
   getFrustrumPlanes,
   mat3,
@@ -68,8 +71,17 @@ import {
 import * as matrix from "#src/util/matrix.js";
 import { GLBuffer } from "#src/webgl/buffer.js";
 import type { GL } from "#src/webgl/context.js";
-import { parameterizedEmitterDependentShaderGetter } from "#src/webgl/dynamic_shader.js";
+import {
+  parameterizedEmitterDependentShaderGetter,
+  shaderCodeWithLineDirective,
+  WatchableShaderError,
+} from "#src/webgl/dynamic_shader.js";
 import type { ShaderBuilder, ShaderProgram } from "#src/webgl/shader.js";
+import {
+  addControlsToBuilder,
+  setControlsInShader,
+  ShaderControlState,
+} from "#src/webgl/shader_ui_controls.js";
 import type { RPC } from "#src/worker_rpc.js";
 import { registerSharedObjectOwner } from "#src/worker_rpc.js";
 
@@ -349,24 +361,44 @@ export class MeshShaderManager {
   }
 
   makeGetter(layer: RefCounted & { gl: GL; displayState: MeshDisplayState }) {
-    const silhouetteRenderingEnabled = layer.registerDisposer(
+    const parameters = layer.registerDisposer(
       makeCachedDerivedWatchableValue(
-        (x) => x > 0,
-        [layer.displayState.silhouetteRendering],
+        (silhouetteRendering, shaderBuilderState) => {
+          return {
+            silhouetteRenderingEnabled: silhouetteRendering > 0,
+            shaderBuilderState,
+          };
+        },
+        [
+          layer.displayState.silhouetteRendering,
+          layer.displayState.shaderControlState.builderState,
+        ],
       ),
     );
+
     return parameterizedEmitterDependentShaderGetter(layer, layer.gl, {
       memoizeKey: `mesh/MeshShaderManager/${this.fragmentRelativeVertices}/${this.vertexPositionFormat}`,
-      parameters: silhouetteRenderingEnabled,
-      defineShader: (builder, silhouetteRenderingEnabled) => {
+      parameters,
+      encodeParameters: (p) => {
+        return `${p.silhouetteRenderingEnabled}/${p.shaderBuilderState.parseResult.code}`;
+      },
+      shaderError: layer.displayState.shaderError,
+      defineShader: (
+        builder,
+        { silhouetteRenderingEnabled, shaderBuilderState },
+      ) => {
+        addControlsToBuilder(shaderBuilderState, builder);
         this.vertexPositionHandler.defineShader(builder);
         builder.addAttribute("highp vec2", "aVertexNormal");
         builder.addVarying("highp vec4", "vColor");
+        builder.addVarying("highp float", "vLightingFactor"); // TODO do I keep this here?
         builder.addUniform("highp vec4", "uLightDirection");
         builder.addUniform("highp vec4", "uColor");
         builder.addUniform("highp mat3", "uNormalMatrix");
         builder.addUniform("highp mat4", "uModelViewProjection");
         builder.addUniform("highp uint", "uPickID");
+        builder.addUniform("highp uvec2", "uID");
+
         if (silhouetteRenderingEnabled) {
           builder.addUniform("highp float", "uSilhouettePower");
         }
@@ -375,6 +407,12 @@ export class MeshShaderManager {
           builder.addUniform("highp vec3", "uFragmentShape");
         }
         builder.addVertexCode(glsl_decodeNormalOctahedronSnorm8);
+
+        builder.addVertexCode(`
+          vec4 defaultColor() { return uColor; }
+          void setColor(vec4 color) {
+            vColor = color;
+          }`);
         let vertexMain = "";
         if (this.fragmentRelativeVertices) {
           vertexMain += `
@@ -393,22 +431,33 @@ vec3 origNormal = decodeNormalOctahedronSnorm8(aVertexNormal);
 vec3 normal = normalize(uNormalMatrix * (normalMultiplier * origNormal));
 float absCosAngle = abs(dot(normal, uLightDirection.xyz));
 float lightingFactor = absCosAngle + uLightDirection.w;
-vColor = vec4(lightingFactor * uColor.rgb, uColor.a);
+vColor = uColor;
+vColor = segmentColor(vColor);
+vColor = vec4(lightingFactor * vColor.rgb, vColor.a);
+vLightingFactor = lightingFactor;
 `;
         if (silhouetteRenderingEnabled) {
           vertexMain += `
 vColor *= pow(1.0 - absCosAngle, uSilhouettePower);
 `;
         }
         builder.setVertexMain(vertexMain);
+        const shaderManager = new SegmentColorShaderManager("getColor");
+        shaderManager.defineShader(builder);
         builder.setFragmentMain("emit(vColor, uPickID);");
+        const segmentColor = shaderCodeWithLineDirective(
+          shaderBuilderState.parseResult.code,
+        );
+        builder.addVertexCode(segmentColor + "\n");
       },
     });
   }
 }
 
 export interface MeshDisplayState extends SegmentationDisplayState3D {
   silhouetteRendering: WatchableValueInterface<number>;
+  shaderControlState: ShaderControlState;
+  shaderError: WatchableShaderError;
 }
 
 export class MeshLayer extends PerspectiveViewRenderLayer<ThreeDimensionalRenderLayerAttachmentState> {
@@ -432,6 +481,14 @@ export class MeshLayer extends PerspectiveViewRenderLayer<ThreeDimensionalRender
     this.registerDisposer(
       displayState.silhouetteRendering.changed.add(this.redrawNeeded.dispatch),
     );
+    this.registerDisposer(
+      displayState.shaderControlState.changed.add(this.redrawNeeded.dispatch),
+    );
+    this.registerDisposer(
+      displayState.shaderControlState.parseResult.changed.add(
+        this.redrawNeeded.dispatch,
+      ),
+    );
 
     const sharedObject = (this.backend = this.registerDisposer(
       new SegmentationLayerSharedObject(
@@ -490,11 +547,19 @@ export class MeshLayer extends PerspectiveViewRenderLayer<ThreeDimensionalRender
     if (modelMatrix === undefined) {
       return;
     }
-    const { shader } = this.getShader(renderContext.emitter);
+    const { shader, parameters } = this.getShader(renderContext.emitter);
     if (shader === null) return;
     shader.bind();
     meshShaderManager.beginLayer(gl, shader, renderContext, this.displayState);
     meshShaderManager.beginModel(gl, shader, renderContext, modelMatrix);
+    setControlsInShader;
+    this.displayState.silhouetteRendering;
+    setControlsInShader(
+      gl,
+      shader,
+      this.displayState.shaderControlState,
+      parameters.shaderBuilderState.parseResult.controls,
+    );
 
     const manifestChunks = this.source.chunks;
 
@@ -774,6 +839,8 @@ export class MultiscaleMeshLayer extends PerspectiveViewRenderLayer<ThreeDimensi
       displayState.silhouetteRendering.changed.add(this.redrawNeeded.dispatch),
     );
 
+    console.log("are we multiscale?");
+
     const sharedObject = (this.backend = this.registerDisposer(
       new SegmentationLayerSharedObject(
         chunkManager,
@@ -829,7 +896,9 @@ export class MultiscaleMeshLayer extends PerspectiveViewRenderLayer<ThreeDimensi
       attachment,
     );
     if (modelMatrix === undefined) return;
-    const { shader } = this.getShader(renderContext.emitter);
+    console.log("get shader multiscale!");
+    const { shader, parameters } = this.getShader(renderContext.emitter);
+    parameters;
     if (shader === null) return;
     shader.bind();
     meshShaderManager.beginLayer(gl, shader, renderContext, this.displayState);
@@ -900,6 +969,7 @@ export class MultiscaleMeshLayer extends PerspectiveViewRenderLayer<ThreeDimensi
         if (renderContext.emitPickID) {
           meshShaderManager.setPickID(gl, shader, pickIndex!);
         }
+
         getMultiscaleChunksToDraw(
           manifest,
           modelViewProjection,