segment color shader

Author: chrisj

src/layer/annotation/index.ts

@@ -770,7 +770,7 @@ export class AnnotationUserLayer extends Base {
 function makeShaderCodeWidget(layer: AnnotationUserLayer) {
   return new ShaderCodeWidget({
     shaderError: layer.annotationDisplayState.shaderError,
-    fragmentMain: layer.annotationDisplayState.shader,
+    fragment: layer.annotationDisplayState.shader,
     shaderControlState: layer.annotationDisplayState.shaderControls,
   });
 }

src/layer/image/index.ts

@@ -461,7 +461,7 @@ export class ImageUserLayer extends Base {
 function makeShaderCodeWidget(layer: ImageUserLayer) {
   return new ShaderCodeWidget({
     shaderError: layer.shaderError,
-    fragmentMain: layer.fragmentMain,
+    fragment: layer.fragmentMain,
     shaderControlState: layer.shaderControlState,
   });
 }

src/layer/segmentation/index.ts

@@ -129,7 +129,14 @@ import {
   verifyString,
 } from "#src/util/json.js";
 import { Signal } from "#src/util/signal.js";
-import { makeWatchableShaderError } from "#src/webgl/dynamic_shader.js";
+import { GLBuffer } from "#src/webgl/buffer.js";
+import { initializeWebGL } from "#src/webgl/context.js";
+import {
+  makeTrackableFragmentMain,
+  makeWatchableShaderError,
+  parameterizedEmitterDependentShaderGetter,
+} from "#src/webgl/dynamic_shader.js";
+import { ShaderControlState } from "#src/webgl/shader_ui_controls.js";
 import type { DependentViewContext } from "#src/widget/dependent_view_widget.js";
 import { registerLayerShaderControlsTool } from "#src/widget/shader_controls.js";
 
@@ -410,12 +417,22 @@ class LinkedSegmentationGroupState<
   }
 }
 
+const DEFAULT_FRAGMENT_SEGMENT_COLOR = `
+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();
+  fragmentSegmentColor = makeTrackableFragmentMain(DEFAULT_FRAGMENT_SEGMENT_COLOR);
+  segmentColorShaderControlState = new ShaderControlState(this.fragmentSegmentColor);
   renderScaleHistogram = new RenderScaleHistogram();
   renderScaleTarget = trackableRenderScaleTarget(1);
   selectSegment: (id: bigint, pin: boolean | "toggle") => void;
@@ -551,6 +570,65 @@ class SegmentationUserLayerDisplayState implements SegmentationDisplayState {
     this.layer.moveToSegment(id);
   };
 
+  makeSegmentColorShaderGetter = () => {
+    const gl = initializeWebGL(new OffscreenCanvas(1, 1));
+    const parameters = this.segmentColorShaderControlState.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 +767,9 @@ export class SegmentationUserLayer extends Base {
     this.displayState.linkedSegmentationGroup.changed.add(() =>
       this.updateDataSubsourceActivations(),
     );
+    this.displayState.fragmentSegmentColor.changed.add(
+      this.specificationChanged.dispatch,
+    );
     this.tabs.add("rendering", {
       label: "Render",
       order: -100,
@@ -996,6 +1077,9 @@ export class SegmentationUserLayer extends Base {
     this.displayState.ignoreNullVisibleSet.restoreState(
       specification[json_keys.IGNORE_NULL_VISIBLE_SET_JSON_KEY],
     );
+    this.displayState.fragmentSegmentColor.restoreState(
+      specification[json_keys.SHADER_JSON_KEY],
+    );
 
     const { skeletonRenderingOptions } = this.displayState;
     skeletonRenderingOptions.restoreState(
@@ -1066,6 +1150,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.fragmentSegmentColor.toJSON();
 
     const { linkedSegmentationGroup, linkedSegmentationColorGroup } =
       this.displayState;
@@ -1365,6 +1450,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 +1487,10 @@ registerLayerTypeDetector((subsource) => {
   return undefined;
 });
 
+registerLayerShaderControlsTool(SegmentationUserLayer, (layer) => ({
+  shaderControlState: layer.displayState.segmentColorShaderControlState,
+}));
+
 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/layer/single_mesh/index.ts

@@ -47,7 +47,7 @@ import {
 } from "#src/widget/shader_controls.js";
 import { Tab } from "#src/widget/tab_view.js";
 
-const SHADER_JSON_KEY = "shader";
+const SEGMENT_COLOR_SHADER_JSON_KEY = "segmentColorShader";
 const SHADER_CONTROLS_JSON_KEY = "shaderControls";
 const CODE_VISIBLE_KEY = "codeVisible";
 
@@ -82,7 +82,7 @@ export class SingleMeshUserLayer extends UserLayer {
   restoreState(specification: any) {
     super.restoreState(specification);
     this.codeVisible.restoreState(specification[CODE_VISIBLE_KEY]);
-    this.displayState.fragmentMain.restoreState(specification[SHADER_JSON_KEY]);
+    this.displayState.fragmentMain.restoreState(specification[SEGMENT_COLOR_SHADER_JSON_KEY]);
     this.displayState.shaderControlState.restoreState(
       specification[SHADER_CONTROLS_JSON_KEY],
     );
@@ -121,7 +121,7 @@ export class SingleMeshUserLayer extends UserLayer {
 
   toJSON() {
     const x = super.toJSON();
-    x[SHADER_JSON_KEY] = this.displayState.fragmentMain.toJSON();
+    x[SEGMENT_COLOR_SHADER_JSON_KEY] = this.displayState.fragmentMain.toJSON();
     x[SHADER_CONTROLS_JSON_KEY] = this.displayState.shaderControlState.toJSON();
     x[CODE_VISIBLE_KEY] = this.codeVisible.toJSON();
     return x;
@@ -133,7 +133,7 @@ export class SingleMeshUserLayer extends UserLayer {
 
 function makeShaderCodeWidget(layer: SingleMeshUserLayer) {
   return new ShaderCodeWidget({
-    fragmentMain: layer.displayState.fragmentMain,
+    fragment: layer.displayState.fragmentMain,
     shaderError: layer.displayState.shaderError,
     shaderControlState: layer.displayState.shaderControlState,
   });

src/mesh/frontend.ts

@@ -43,6 +43,7 @@ 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,
@@ -68,8 +69,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 type { WatchableShaderError } from "#src/webgl/dynamic_shader.js";
+import {
+  parameterizedEmitterDependentShaderGetter,
+  shaderCodeWithLineDirective,
+} from "#src/webgl/dynamic_shader.js";
 import type { ShaderBuilder, ShaderProgram } from "#src/webgl/shader.js";
+import type { ShaderControlState } from "#src/webgl/shader_ui_controls.js";
+import {
+  addControlsToBuilder,
+  setControlsInShader,
+} from "#src/webgl/shader_ui_controls.js";
 import type { RPC } from "#src/worker_rpc.js";
 import { registerSharedObjectOwner } from "#src/worker_rpc.js";
 
@@ -349,16 +359,33 @@ 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.segmentColorShaderControlState.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");
@@ -367,6 +394,7 @@ export class MeshShaderManager {
         builder.addUniform("highp mat3", "uNormalMatrix");
         builder.addUniform("highp mat4", "uModelViewProjection");
         builder.addUniform("highp uint", "uPickID");
+
         if (silhouetteRenderingEnabled) {
           builder.addUniform("highp float", "uSilhouettePower");
         }
@@ -393,22 +421,32 @@ 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);
 `;
         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>;
+  segmentColorShaderControlState: ShaderControlState;
+  shaderError: WatchableShaderError;
 }
 
 export class MeshLayer extends PerspectiveViewRenderLayer<ThreeDimensionalRenderLayerAttachmentState> {
@@ -432,6 +470,14 @@ export class MeshLayer extends PerspectiveViewRenderLayer<ThreeDimensionalRender
     this.registerDisposer(
       displayState.silhouetteRendering.changed.add(this.redrawNeeded.dispatch),
     );
+    this.registerDisposer(
+      displayState.segmentColorShaderControlState.changed.add(this.redrawNeeded.dispatch),
+    );
+    this.registerDisposer(
+      displayState.segmentColorShaderControlState.parseResult.changed.add(
+        this.redrawNeeded.dispatch,
+      ),
+    );
 
     const sharedObject = (this.backend = this.registerDisposer(
       new SegmentationLayerSharedObject(
@@ -490,11 +536,17 @@ 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(
+      gl,
+      shader,
+      this.displayState.segmentColorShaderControlState,
+      parameters.shaderBuilderState.parseResult.controls,
+    );
 
     const manifestChunks = this.source.chunks;