@vitessce/neuroglancer 3.9.6 → 3.9.8

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "@vitessce/neuroglancer",
-  "version": "3.9.6",
+  "version": "3.9.8",
   "author": "Gehlenborg Lab",
   "homepage": "http://vitessce.io",
   "repository": {
@@ -20,14 +20,14 @@
     "lodash-es": "^4.17.21",
     "three": "^0.154.0",
     "react": "18.3.1",
-    "@vitessce/neuroglancer-workers": "3.9.6",
-    "@vitessce/vit-s": "3.9.6",
-    "@vitessce/styles": "3.9.6",
-    "@vitessce/utils": "3.9.6",
-    "@vitessce/constants-internal": "3.9.6",
-    "@vitessce/sets-utils": "3.9.6",
-    "@vitessce/tooltip": "3.9.6",
-    "@vitessce/legend": "3.9.6"
+    "@vitessce/neuroglancer-workers": "3.9.8",
+    "@vitessce/vit-s": "3.9.8",
+    "@vitessce/styles": "3.9.8",
+    "@vitessce/sets-utils": "3.9.8",
+    "@vitessce/constants-internal": "3.9.8",
+    "@vitessce/utils": "3.9.8",
+    "@vitessce/tooltip": "3.9.8",
+    "@vitessce/legend": "3.9.8"
   },
   "devDependencies": {
     "@testing-library/jest-dom": "^6.6.3",

package/src/Neuroglancer.js

@@ -1,12 +1,14 @@
 /* eslint-disable react-refresh/only-export-components */
 import React, { PureComponent, Suspense } from 'react';
-import { ChunkWorker } from '@vitessce/neuroglancer-workers';
+import { ChunkWorker, AsyncComputationWorker } from '@vitessce/neuroglancer-workers';
 import { NeuroglancerGlobalStyles } from './styles.js';
 
 const LazyReactNeuroglancer = React.lazy(() => import('./ReactNeuroglancer.js'));
 
 function createWorker() {
-  return new ChunkWorker();
+  const worker = new ChunkWorker();
+  worker.AsyncComputationWorker = AsyncComputationWorker;
+  return worker;
 }
 export class NeuroglancerComp extends PureComponent {
   constructor(props) {
@@ -14,8 +16,6 @@ export class NeuroglancerComp extends PureComponent {
     this.bundleRoot = createWorker();
     this.cellColorMapping = props.cellColorMapping;
     this.justReceivedExternalUpdate = false;
-    this.prevElement = null;
-    this.prevClickHandler = null;
     this.prevMouseStateChanged = null;
     this.prevHoverHandler = null;
     this.onViewerStateChanged = this.onViewerStateChanged.bind(this);
@@ -32,41 +32,49 @@ export class NeuroglancerComp extends PureComponent {
     if (viewerRef) {
       // Mount
       const { viewer } = viewerRef;
-      this.prevElement = viewer.element;
       this.prevMouseStateChanged = viewer.mouseState.changed;
-      viewer.inputEventBindings.sliceView.set('at:dblclick0', () => {});
+      // For now, can omit the sliceView bindings, as we only use perspectiveView
+      // viewer.inputEventBindings.sliceView.set('at:dblclick0', () => {});
       viewer.inputEventBindings.perspectiveView.set('at:dblclick0', () => {});
-      this.prevClickHandler = (event) => {
-        if (event.button === 0) {
-          // Wait for mouseState to update
-          requestAnimationFrame(() => {
-            const { pickedValue, pickedRenderLayer } = viewer.mouseState;
-            // Only trigger selection when a segment is clicked rather than any click on the view
-            if (pickedValue && pickedValue.low !== undefined && pickedRenderLayer) {
-              this.latestOnSegmentClick?.(pickedValue.low);
+
+      // Disable space interaction to prevent triggering 4panels layout.
+      viewer.inputEventBindings.sliceView.set('at:space', () => {});
+      viewer.inputEventBindings.perspectiveView.set('at:space', () => {});
+
+      // Remap plain wheel to  ctrl+wheel (zoom) action
+      // by traversing the parent binding maps.
+      const remapWheelToZoom = (map) => {
+        if (map.bindings) {
+          const ctrlWheelAction = map.bindings.get('at:control+wheel');
+          if (ctrlWheelAction) {
+            // Replace plain wheel with the zoom action
+            map.bindings.set('at:wheel', ctrlWheelAction);
+            const ctrlWheelBubble = map.bindings.get('bubble:control+wheel');
+            if (ctrlWheelBubble) {
+              map.bindings.set('bubble:wheel', ctrlWheelBubble);
             }
-          });
+          }
+        }
+        if (map.parents) {
+          map.parents.forEach(p => remapWheelToZoom(p));
         }
       };
+
+      remapWheelToZoom(viewer.inputEventBindings.perspectiveView);
+
       this.prevHoverHandler = () => {
         if (viewer.mouseState.pickedValue !== undefined) {
           const pickedSegment = viewer.mouseState.pickedValue;
           this.latestOnSelectHoveredCoords?.(pickedSegment?.low);
         }
       };
-      viewer.element.addEventListener('mouseup', this.prevClickHandler);
+
       viewer.mouseState.changed.add(this.prevHoverHandler);
     } else {
-      // Unmount (viewerRef is null)
-      if (this.prevElement && this.prevClickHandler) {
-        this.prevElement.removeEventListener('mouseup', this.prevClickHandler);
-        this.prevClickHandler = null;
-      }
       if (this.prevMouseStateChanged && this.prevHoverHandler) {
         this.prevMouseStateChanged.remove(this.prevHoverHandler);
         this.prevHoverHandler = null;
       }
-      this.prevElement = null;
       this.prevMouseStateChanged = null;
     }
   }
@@ -77,17 +85,14 @@ export class NeuroglancerComp extends PureComponent {
   }
 
   componentDidUpdate(prevProps) {
-    const { onSegmentClick, onSelectHoveredCoords } = this.props;
-    if (prevProps.onSegmentClick !== onSegmentClick) {
-      this.latestOnSegmentClick = onSegmentClick;
-    }
+    const { onSelectHoveredCoords } = this.props;
     if (prevProps.onSelectHoveredCoords !== onSelectHoveredCoords) {
       this.latestOnSelectHoveredCoords = onSelectHoveredCoords;
     }
   }
 
   render() {
-    const { classes, viewerState, cellColorMapping } = this.props;
+    const { classes, viewerState, cellColorMapping, onLayerLoadingChange } = this.props;
 
     return (
       <>
@@ -98,6 +103,7 @@ export class NeuroglancerComp extends PureComponent {
               brainMapsClientId="NOT_A_VALID_ID"
               viewerState={viewerState}
               onViewerStateChanged={this.onViewerStateChanged}
+              onLayerLoadingChange={onLayerLoadingChange}
               bundleRoot={this.bundleRoot}
               cellColorMapping={cellColorMapping}
               ref={this.onRef}

package/src/NeuroglancerSubscriber.js

@@ -57,6 +57,8 @@ const ROTATION_EPS = 1e-3;
 const TARGET_EPS = 0.5;
 const NG_ROT_COOLDOWN_MS = 120;
 
+const GUIDE_URL = 'https://vitessce.io/docs/ng-guide/';
+
 const LAST_INTERACTION_SOURCE = {
   vitessce: 'vitessce',
   neuroglancer: 'neuroglancer',
@@ -77,6 +79,7 @@ export function NeuroglancerSubscriber(props) {
     downloadButtonVisible,
     removeGridComponent,
     theme,
+    showAxisLines = false,
     title = 'Spatial',
     subtitle = 'Powered by Neuroglancer',
     helpText = ViewHelpMapping.NEUROGLANCER,
@@ -205,6 +208,7 @@ export function NeuroglancerSubscriber(props) {
       CoordinationType.TOOLTIPS_VISIBLE,
       CoordinationType.TOOLTIP_CROSSHAIRS_VISIBLE,
       CoordinationType.LEGEND_VISIBLE,
+      CoordinationType.SPATIAL_POINT_STROKE_WIDTH,
     ],
     coordinationScopes,
     coordinationScopesBy,
@@ -281,13 +285,41 @@ export function NeuroglancerSubscriber(props) {
       segmentationChannelScopesByLayer?.[layerScope]?.forEach((channelScope) => {
         const { obsSets: layerSets, obsIndex: layerIndex } = obsSegmentationsSetsData
           ?.[layerScope]?.[channelScope] || {};
-        if (layerSets && layerIndex) {
-          const {
-            obsSetColor,
-            obsColorEncoding,
-            obsSetSelection,
-            additionalObsSets,
-          } = segmentationChannelCoordination[0][layerScope][channelScope];
+        const {
+          obsSetColor,
+          obsColorEncoding,
+          obsSetSelection,
+          additionalObsSets,
+          spatialChannelColor,
+          spatialChannelOpacity,
+        } = segmentationChannelCoordination[0][layerScope][channelScope];
+        if (obsColorEncoding === 'spatialChannelColor') {
+          // All segments get the same static channel color
+          if (layerIndex && spatialChannelColor) {
+            const hex = rgbToHex(spatialChannelColor);
+            const ngCellColors = {};
+
+            if (obsSetSelection?.length > 0) {
+              // Only color the segments belonging to selected sets.
+              const mergedCellSets = mergeObsSets(layerSets, additionalObsSets);
+              const selectedIds = new Set();
+              obsSetSelection.forEach((setPath) => {
+                const rootNode = mergedCellSets?.tree?.find(n => n.name === setPath[0]);
+                const leafNode = setPath.length > 1
+                  ? rootNode?.children?.find(n => n.name === setPath[1])
+                  : rootNode;
+                leafNode?.set?.forEach(([id]) => selectedIds.add(String(id)));
+              });
+              layerIndex.forEach((id) => {
+                if (selectedIds.has(String(id))) {
+                  ngCellColors[id] = hex;
+                }
+              });
+            }
+            result[layerScope][channelScope] = ngCellColors;
+            result[layerScope].opacity = spatialChannelOpacity ?? 1.0;
+          }
+        } else if (layerSets && layerIndex) {
           const mergedCellSets = mergeObsSets(layerSets, additionalObsSets);
           const cellColors = getCellColors({
             cellSets: mergedCellSets,
@@ -301,14 +333,8 @@ export function NeuroglancerSubscriber(props) {
           cellColors.forEach((color, i) => {
             ngCellColors[i] = rgbToHex(color);
           });
-          /* // TODO: Is this necessary?
-          const obsColorIndices = treeToCellSetColorIndicesBySetNames(
-            mergedLayerSets,
-            obsSetSelection,
-            obsSetColor,
-          );
-          */
           result[layerScope][channelScope] = ngCellColors;
+          result[layerScope].opacity = spatialChannelOpacity ?? 1.0;
         }
       });
     });
@@ -327,6 +353,7 @@ export function NeuroglancerSubscriber(props) {
   // Obtain the Neuroglancer viewerState object.
   const initalViewerState = useNeuroglancerViewerState(
     theme,
+    showAxisLines,
     segmentationLayerScopes,
     segmentationChannelScopesByLayer,
     segmentationLayerCoordination,
@@ -381,6 +408,9 @@ export function NeuroglancerSubscriber(props) {
     orbit: spatialRotationOrbit,
   });
 
+  // Track layer loading state for showing loading indicator
+  const [isLayersLoaded, setIsLayersLoaded] = useState(false);
+
   // Track the last coord values we saw, and only mark "vitessce"
   // when *those* actually change. This prevents cell set renders
   // from spoofing the source.
@@ -512,6 +542,8 @@ export function NeuroglancerSubscriber(props) {
   }, []);
 
   const onSegmentClick = useCallback((value) => {
+    // Note: this callback is no longer called by the child component.
+    // Reference: https://github.com/vitessce/vitessce/pull/2439
     if (value) {
       const id = String(value);
       const selectedCellIds = [id];
@@ -536,14 +568,19 @@ export function NeuroglancerSubscriber(props) {
     setCellColorEncoding, setCellSetColor, setCellSetSelection,
   ]);
 
-  // Get the ultimate cellColorMapping to pass to NeuroglancerComp as a prop.
-  // For now, we take the first layer and channel for cell colors.
-  const cellColorMapping = useMemo(() => (segmentationColorMapping
-    ?.[segmentationLayerScopes?.[0]]
-    ?.[segmentationChannelScopesByLayer?.[segmentationLayerScopes?.[0]]?.[0]]
-    ?? {}
-  ), [segmentationColorMapping]);
+  // Get the ultimate cellColorMapping for each layer to pass to NeuroglancerComp as a prop.
 
+  const cellColorMappingByLayer = useMemo(() => {
+    const result = {};
+    segmentationLayerScopes?.forEach((layerScope) => {
+      const channelScope = segmentationChannelScopesByLayer?.[layerScope]?.[0];
+      result[layerScope] = {
+        colors: segmentationColorMapping?.[layerScope]?.[channelScope] ?? {},
+        opacity: segmentationColorMapping?.[layerScope]?.opacity ?? 1.0,
+      };
+    });
+    return result;
+  }, [segmentationColorMapping, segmentationLayerScopes, segmentationChannelScopesByLayer]);
 
   // TODO: try to simplify using useMemoCustomComparison?
   // This would allow us to refactor a lot of the checking-for-changes logic into a comparison function,
@@ -556,9 +593,6 @@ export function NeuroglancerSubscriber(props) {
       return current;
     }
 
-    const nextSegments = Object.keys(cellColorMapping);
-    const prevLayer = current?.layers?.[0] || {};
-    const prevSegments = prevLayer.segments || [];
     const { projectionScale, projectionOrientation, position } = current;
 
     // Did Vitessce coords change vs the *previous* render?
@@ -695,20 +729,24 @@ export function NeuroglancerSubscriber(props) {
       lastInteractionSource.current = null;
     }
 
-    const newLayer0 = {
-      ...prevLayer,
-      segments: nextSegments,
-      segmentColors: cellColorMapping,
-    };
-
+    const updatedLayers = current?.layers?.map((layer, idx) => {
+      const layerScope = segmentationLayerScopes?.[idx];
+      const layerColorMapping = cellColorMappingByLayer?.[layerScope]?.colors ?? {};
+      const layerSegments = Object.keys(layerColorMapping);
+      return {
+        ...layer,
+        segments: layerSegments,
+        segmentColors: layerColorMapping,
+        objectAlpha: cellColorMappingByLayer?.[layerScope]?.opacity ?? 1.0,
+      };
+    }) ?? [];
 
     const updated = {
       ...current,
       projectionScale: nextProjectionScale,
       projectionOrientation: nextOrientation,
       position: nextPosition,
-      layers: prevSegments.length === 0 ? [newLayer0, ...(current?.layers?.slice(1)
-        || [])] : current?.layers,
+      layers: updatedLayers,
     };
 
     latestViewerStateRef.current = updated;
@@ -724,7 +762,7 @@ export function NeuroglancerSubscriber(props) {
     };
 
     return updated;
-  }, [cellColorMapping, spatialZoom, spatialRotationX, spatialRotationY,
+  }, [cellColorMappingByLayer, spatialZoom, spatialRotationX, spatialRotationY,
     spatialRotationZ, spatialTargetX, spatialTargetY, initalViewerState,
     latestViewerStateIteration]);
 
@@ -732,6 +770,10 @@ export function NeuroglancerSubscriber(props) {
     setCellHighlight(String(obsId));
   }, [setCellHighlight]);
 
+  const handleLayerLoadingChange = useCallback((isLoaded) => {
+    setIsLayersLoaded(isLoaded);
+  }, []);
+
   // TODO: if all cells are deselected, a black view is shown, rather we want to show empty NG view?
   // if (!cellColorMapping || Object.keys(cellColorMapping).length === 0) {
   //   return;
@@ -741,6 +783,7 @@ export function NeuroglancerSubscriber(props) {
   // console.log(derivedViewerState);
 
   return (
+
     <TitleInfo
       title={title}
       info={subtitle}
@@ -750,33 +793,43 @@ export function NeuroglancerSubscriber(props) {
       closeButtonVisible={closeButtonVisible}
       downloadButtonVisible={downloadButtonVisible}
       removeGridComponent={removeGridComponent}
-      isReady={isReady}
+      isReady={isReady && isLayersLoaded}
       errors={errors}
       withPadding={false}
+      guideUrl={GUIDE_URL}
     >
-      <div style={{ position: 'relative', width: '100%', height: '100%' }} ref={containerRef}>
-        <div style={{ position: 'absolute', top: 0, right: 0, zIndex: 50 }}>
-          <MultiLegend
-            theme="dark"
-            maxHeight={ngHeight}
-            segmentationLayerScopes={segmentationLayerScopes}
-            segmentationLayerCoordination={segmentationLayerCoordination}
-            segmentationChannelScopesByLayer={segmentationChannelScopesByLayer}
-            segmentationChannelCoordination={segmentationChannelCoordination}
-          />
-        </div>
+      {hasLayers ? (
+        <div style={{ position: 'relative', width: '100%', height: '100%' }} ref={containerRef}>
+          <div style={{ position: 'absolute', top: 0, right: 0, zIndex: 50 }}>
+            <MultiLegend
+              theme="dark"
+              maxHeight={ngHeight}
+
+              // Segmentations
+              segmentationLayerScopes={segmentationLayerScopes}
+              segmentationLayerCoordination={segmentationLayerCoordination}
+              segmentationChannelScopesByLayer={segmentationChannelScopesByLayer}
+              segmentationChannelCoordination={segmentationChannelCoordination}
+
+              // Points
+              pointLayerScopes={pointLayerScopes}
+              pointLayerCoordination={pointLayerCoordination}
+              pointMultiIndicesData={pointMultiIndicesData}
+            />
+          </div>
 
-        {hasLayers ? (
           <NeuroglancerComp
             classes={classes}
             onSegmentClick={onSegmentClick}
             onSelectHoveredCoords={onSegmentHighlight}
             viewerState={derivedViewerState}
-            cellColorMapping={cellColorMapping}
+            cellColorMapping={cellColorMappingByLayer}
             setViewerState={handleStateUpdate}
+            onLayerLoadingChange={handleLayerLoadingChange}
           />
-        ) : null}
-      </div>
+        </div>
+      ) : null}
     </TitleInfo>
+
   );
 }

package/src/ReactNeuroglancer.js

@@ -67,6 +67,10 @@ let viewerNoKey;
  * @property {() => void} onSelectionDetailsStateChanged
  * A function of the form `() => {}` to respond to selection changes in the viewer.
  * @property {() => void} onViewerStateChanged
+ * @property {(isLoaded: boolean) => void} onLayerLoadingChange
+ * A function of the form `(isLoaded) => {}`, called when layer loading state changes.
+ * The `isLoaded` argument will be `true` when all segmentation layers have finished loading
+ * their data sources, or `false` when layers are still loading.
  *
  * @property {Array<Object>} callbacks
  * // ngServer: string,
@@ -411,6 +415,7 @@ export default class Neuroglancer extends React.Component {
     onVisibleChanged: null,
     onSelectionDetailsStateChanged: null,
     onViewerStateChanged: null,
+    onLayerLoadingChange: null,
     key: null,
     callbacks: [],
     ngServer: 'https://neuroglancer-demo.appspot.com/',
@@ -427,7 +432,7 @@ export default class Neuroglancer extends React.Component {
     this.disposers = [];
     this.prevColorOverrides = new Set();
     this.overrideColorsById = Object.create(null);
-    this.allKnownIds = new Set();
+    this.allKnownIdsByLayer = {};
   }
 
   minimalPoseSnapshot = () => {
@@ -476,30 +481,36 @@ export default class Neuroglancer extends React.Component {
   };
 
   /* To add colors to the segments, turning unselected to grey  */
-  applyColorsAndVisibility = (cellColorMapping) => {
+  applyColorsAndVisibility = (cellColorMappingByLayer) => {
     if (!this.viewer) return;
-    // Track all ids we've ever seen so we can grey the ones
-    // that drop out of the current selection.
-    const selected = { ...(cellColorMapping || {}) }; // clone, don't mutate props
-    for (const id of Object.keys(selected)) this.allKnownIds.add(id);
-    // If empty on first call, seed from initial segmentColors (if present)
-    if (this.allKnownIds.size === 0) {
-      const init = this.props.viewerState?.layers?.[0]?.segmentColors || {};
-      for (const id of Object.keys(init)) this.allKnownIds.add(id);
-    }
-
-    // Build a full color table: selected keep their hex, others grey
-    const fullSegmentColors = {};
-    for (const id of this.allKnownIds) {
-      fullSegmentColors[id] = selected[id] || GREY_HEX;
-    }
-    // Patch layers with the new segmentColors (pose untouched)
+    // Build full color table per layer
     const baseLayers = (this.props.viewerState?.layers)
       ?? (this.viewer.state.toJSON().layers || []);
 
-    const newLayers = baseLayers.map((layer, idx) => {
-      // if only one layer, take that or check layer.type === 'segmentation'.
-      if (idx === 0 || layer?.type === 'segmentation') {
+    const newLayers = baseLayers.map((layer) => {
+      // Match layerScope by checking if the NG layer name contains the scope key.
+      // NG layer names are of the form:
+      // "obsSegmentations-init_A_obsSegmentations_0-init_A_obsSegmentations_0"
+      const layerScope = Object.keys(cellColorMappingByLayer).find(scope => layer.name?.includes(scope));
+
+      const selected = { ...(cellColorMappingByLayer[layerScope]?.colors || {}) };
+
+      // Track all known IDs for this layer scope
+      if (!this.allKnownIdsByLayer) this.allKnownIdsByLayer = {};
+      if (!this.allKnownIdsByLayer[layerScope]) {
+        this.allKnownIdsByLayer[layerScope] = new Set();
+      }
+      for (const id of Object.keys(selected)) {
+        this.allKnownIdsByLayer[layerScope].add(id);
+      }
+
+      // Build a full color table: selected keep their hex, others grey
+      const fullSegmentColors = {};
+      for (const id of this.allKnownIdsByLayer[layerScope] || []) {
+        fullSegmentColors[id] = selected[id] || GREY_HEX;
+      }
+
+      if (layer.type === 'segmentation') {
         return { ...layer, segmentColors: fullSegmentColors };
       }
       return layer;
@@ -619,12 +630,46 @@ export default class Neuroglancer extends React.Component {
       viewerNoKey = this.viewer;
     }
 
-    // TODO: This is purely for debugging and we need to remove it.
+    const { visibleChunksChanged } = this.viewer.chunkQueueManager;
+    let firstChunkLoaded = false;
+    this.disposers.push(visibleChunksChanged.add(() => {
+      if (!firstChunkLoaded) {
+        for (const layer of this.viewer.layerManager.managedLayers) {
+          if (layer.layer instanceof SegmentationUserLayer) {
+            const hasVisibleChunk = layer.layer.renderLayers?.some((rl) => {
+              const {
+                numVisibleChunksAvailable,
+                numVisibleChunksNeeded,
+              } = rl.layerChunkProgressInfo || {};
+              if (!numVisibleChunksNeeded || !numVisibleChunksAvailable) return false;
+              // Neuroglancer only shows chunks when a certain % is loaded.
+              // The 0.25 is from testing different values, can be reduced to 0.2 to shorten loader time
+              return (numVisibleChunksAvailable / numVisibleChunksNeeded) > 0.25;
+            });
+            if (hasVisibleChunk) {
+              firstChunkLoaded = true;
+              // Two frames to avoid flash while the following two happens
+              // Neuroglancer issues WebGL draw calls
+              requestAnimationFrame(() => {
+              // GPU has painted, pixels visible on screen
+                requestAnimationFrame(() => {
+                  this.props.onLayerLoadingChange?.(true);
+                });
+              });
+              return;
+            }
+          }
+        }
+      }
+    }));
+    this.disposers.push(() => { firstChunkLoaded = false; });
+
+    // TODO: This is purely for debugging - exposes the NG viewer to be tested via console
     // window.viewer = this.viewer;
   }
 
   componentDidUpdate(prevProps, prevState) {
-    const { viewerState, cellColorMapping } = this.props;
+    const { viewerState, cellColorMapping: cellColorMappingByLayer } = this.props;
     // The restoreState() call clears the 'selected' (hovered on) segment, which is needed
     // by Neuroglancer's code to toggle segment visibilty on a mouse click.  To free the user
     // from having to move the mouse before clicking, save the selected segment and restore
@@ -661,6 +706,27 @@ export default class Neuroglancer extends React.Component {
       if (layer.layer instanceof SegmentationUserLayer) {
         const { segmentSelectionState } = layer.layer.displayState;
         segmentSelectionState.set(selectedSegments[layer.name]);
+        const layerScope = Object.keys(cellColorMappingByLayer).find(
+          scope => layer.name?.includes(scope),
+        );
+        if (layerScope) {
+          const opacity = cellColorMappingByLayer[layerScope]?.opacity ?? 1.0;
+          layer.layer.displayState.objectAlpha.value = opacity;
+        }
+      }
+      // Update annotation layer shaders from viewerState config,
+      // skipping update if shader is unchanged to avoid costly re-renders
+      if (layer.layer instanceof AnnotationUserLayer) {
+        const matchingLayer = (viewerState?.layers || []).find(
+          l => l.name === layer.name,
+        );
+        if (matchingLayer?.shader) {
+          /* eslint-disable-next-line no-underscore-dangle */
+          const currentShader = layer.layer.annotationDisplayState.shader.value_;
+          if (currentShader !== matchingLayer.shader) {
+            layer.layer.annotationDisplayState.shader.value = matchingLayer.shader;
+          }
+        }
       }
     }
 
@@ -702,8 +768,8 @@ export default class Neuroglancer extends React.Component {
       this.withoutEmitting(() => {
         const layers = Array.isArray(viewerState.layers) ? viewerState.layers : [];
         this.viewer.state.restoreState({ layers });
-        if (cellColorMapping && Object.keys(cellColorMapping).length) {
-          this.applyColorsAndVisibility(cellColorMapping);
+        if (cellColorMappingByLayer && Object.keys(cellColorMappingByLayer).length) {
+          this.applyColorsAndVisibility(cellColorMappingByLayer);
         }
       });
     }
@@ -711,13 +777,16 @@ export default class Neuroglancer extends React.Component {
     // If colors changed (but layers didn’t): re-apply colors
     // this was to avid NG randomly assigning colors to the segments by resetting them
     const prevSize = prevProps.cellColorMapping
-      ? Object.keys(prevProps.cellColorMapping).length : 0;
-    const currSize = cellColorMapping ? Object.keys(cellColorMapping).length : 0;
-    const mappingRefChanged = prevProps.cellColorMapping !== cellColorMapping;
+      ? Object.values(prevProps.cellColorMapping)
+        .reduce((acc, v) => acc + Object.keys(v?.colors || {}).length, 0) : 0;
+    const currSize = cellColorMappingByLayer
+      ? Object.values(cellColorMappingByLayer)
+        .reduce((acc, v) => acc + Object.keys(v?.colors || {}).length, 0) : 0;
+    const mappingRefChanged = prevProps.cellColorMapping !== this.props.cellColorMapping;
     if (!this.didLayersChange(prevVS, viewerState)
       && (mappingRefChanged || prevSize !== currSize)) {
       this.withoutEmitting(() => {
-        this.applyColorsAndVisibility(cellColorMapping);
+        this.applyColorsAndVisibility(cellColorMappingByLayer);
       });
     }
 
@@ -725,7 +794,7 @@ export default class Neuroglancer extends React.Component {
     // We only restore layers (not pose) when sources change OR on the first time segments appear.
     const stripSegFields = layers => (layers || []).map((l) => {
       if (!l) return l;
-      const { segments, segmentColors, ...rest } = l;
+      const { segments, segmentColors, objectAlpha, ...rest } = l;
       return rest; // ignore segments + segmentColors for comparison
     });