From 8efb5fac6a9bb6a4987c8a6fbbab12ad2714dfce Mon Sep 17 00:00:00 2001 From: Alireza Date: Mon, 6 Jul 2026 11:37:44 -0400 Subject: [PATCH 1/5] fix(segmentation): make overlap labelmap strokes fully undoable A stroke that forces a segment move to a private labelmap layer changed three things outside its history memo: the bulk voxel move, the layer registration, and the segment rebinding. Cross-layer overwrite erases were also unrecorded. Undo then restored only the stroke's own voxels, leaving the moved voxels stranded on a layer that should not exist at that point in history. Labelmap memos now carry ordered restore steps (priorSteps/postSteps): moveSegmentToPrivateLabelmap emits a step that reverses/replays the whole move, cross-layer erases are recorded per layer through crossLayerEraseCallback, and a mid-stroke voxel-manager swap rides the committed earlier memo along. Undo walks steps in reverse-chronological order, redo chronologically, so each stroke restores as one unit. --- .../labelmapModel/labelmapEditTransaction.ts | 46 +- .../labelmapModel/privateLabelmap.spec.ts | 436 ++++++++++++++++++ .../labelmapModel/privateLabelmap.ts | 82 +++- .../segmentation/strategies/BrushStrategy.ts | 32 ++ .../strategies/utils/crossLayerErase.ts | 24 + .../segmentation/createLabelmapMemo.spec.ts | 197 ++++++++ .../segmentation/createLabelmapMemo.ts | 91 +++- 7 files changed, 890 insertions(+), 18 deletions(-) create mode 100644 packages/tools/src/stateManagement/segmentation/labelmapModel/privateLabelmap.spec.ts create mode 100644 packages/tools/src/utilities/segmentation/createLabelmapMemo.spec.ts diff --git a/packages/tools/src/stateManagement/segmentation/labelmapModel/labelmapEditTransaction.ts b/packages/tools/src/stateManagement/segmentation/labelmapModel/labelmapEditTransaction.ts index 718b7e45df..6f6e829e26 100644 --- a/packages/tools/src/stateManagement/segmentation/labelmapModel/labelmapEditTransaction.ts +++ b/packages/tools/src/stateManagement/segmentation/labelmapModel/labelmapEditTransaction.ts @@ -13,10 +13,20 @@ import { getSegmentsOnLabelmap, } from './labelmapSegmentBindings'; import { moveSegmentToPrivateLabelmap as defaultMoveSegmentToPrivateLabelmap } from './privateLabelmap'; +import type { LabelmapRestoreStep } from '../../../utilities/segmentation/createLabelmapMemo'; + +/** A batch of cross-layer erases performed by one strategy fill call, with + * everything needed to reverse/replay it on undo/redo. */ +type CrossLayerEraseRecord = { + voxelManager: Types.IVoxelManager; + labelValue: number; + indices: number[]; +}; type MoveSegmentToPrivateLabelmap = ( segmentation: Segmentation, - segmentIndex: number + segmentIndex: number, + options?: { moveStepCallback?: (step: LabelmapRestoreStep) => void } ) => LabelmapLayer | undefined; type BeginLabelmapEditTransactionOptions = { @@ -39,6 +49,10 @@ type LabelmapEditTransaction = { protectedSegmentIndices: number[]; crossLayerEraseBindings: SegmentLabelmapBindingState[]; movedSegment: boolean; + /** Undo/redo step for a segment move performed by this transaction + * (layer registration + bulk voxel move + binding change), for recording + * on the stroke's memo. */ + moveStep?: LabelmapRestoreStep; }; type ResolveLabelmapLayerEditTargetOptions = { @@ -62,6 +76,9 @@ type EraseLabelmapEditTransactionOptions = { isInObject: (point: Types.Point3) => boolean; isInObjectBoundsIJK?: Types.BoundsIJK; imageId?: string; + /** When provided, receives the erased voxels of each touched layer so the + * caller can record them for undo/redo. */ + crossLayerEraseCallback?: (record: CrossLayerEraseRecord) => void; }; function getProtectedSegmentIndicesForLayer( @@ -197,13 +214,16 @@ function beginLabelmapEditTransaction( } ); + let moveStep: LabelmapRestoreStep | undefined; + if (shouldMoveSegment) { const moveSegmentToPrivateLabelmap = options.moveSegmentToPrivateLabelmap ?? defaultMoveSegmentToPrivateLabelmap; const privateLayer = moveSegmentToPrivateLabelmap( segmentation, - segmentIndex + segmentIndex, + { moveStepCallback: (step) => (moveStep = step) } ); const privateBinding = getSegmentBinding(segmentation, segmentIndex); @@ -234,6 +254,7 @@ function beginLabelmapEditTransaction( overwriteSegmentIndices ), movedSegment, + moveStep, }; } @@ -329,6 +350,7 @@ function eraseVolumeLayer( return; } + const erasedIndices: number[] = []; volume.voxelManager.forEach( ({ value, index, pointIJK }) => { if (value !== binding.labelValue) { @@ -343,6 +365,7 @@ function eraseVolumeLayer( } volume.voxelManager.setAtIndex(index, 0); + erasedIndices.push(index); }, { imageData: volume.imageData, @@ -350,6 +373,14 @@ function eraseVolumeLayer( } ); + if (erasedIndices.length) { + options.crossLayerEraseCallback?.({ + voxelManager: volume.voxelManager as Types.IVoxelManager, + labelValue: binding.labelValue, + indices: erasedIndices, + }); + } + volume.voxelManager ?.getArrayOfModifiedSlices?.() ?.forEach((sliceIndex) => modifiedSlices.add(sliceIndex)); @@ -373,6 +404,7 @@ function eraseStackLayer( return; } + const erasedIndices: number[] = []; voxelManager.forEach( ({ value, index, pointIJK }) => { if (value !== binding.labelValue) { @@ -387,6 +419,7 @@ function eraseStackLayer( } voxelManager.setAtIndex(index, 0); + erasedIndices.push(index); }, { imageData: options.referenceImageData, @@ -394,6 +427,14 @@ function eraseStackLayer( } ); + if (erasedIndices.length) { + options.crossLayerEraseCallback?.({ + voxelManager, + labelValue: binding.labelValue, + indices: erasedIndices, + }); + } + const currentSlice = stackViewport.getCurrentImageIdIndex?.(); if (typeof currentSlice === 'number') { modifiedSlices.add(currentSlice); @@ -440,6 +481,7 @@ export { }; export type { BeginLabelmapEditTransactionOptions, + CrossLayerEraseRecord, EraseLabelmapEditTransactionOptions, LabelmapEditTransaction, LabelmapLayerEditTarget, diff --git a/packages/tools/src/stateManagement/segmentation/labelmapModel/privateLabelmap.spec.ts b/packages/tools/src/stateManagement/segmentation/labelmapModel/privateLabelmap.spec.ts new file mode 100644 index 0000000000..585f07b3e3 --- /dev/null +++ b/packages/tools/src/stateManagement/segmentation/labelmapModel/privateLabelmap.spec.ts @@ -0,0 +1,436 @@ +import type { Types } from '@cornerstonejs/core'; +import type vtkImageData from '@kitware/vtk.js/Common/DataModel/ImageData'; +import type { Segmentation } from '../../../types/SegmentationStateTypes'; +import type { LabelmapRestoreStep } from '../../../utilities/segmentation/createLabelmapMemo'; + +jest.mock('@cornerstonejs/core', () => { + const actual = jest.requireActual('@cornerstonejs/core'); + return { + ...actual, + cache: { + getImage: jest.fn(), + getVolume: jest.fn(), + removeVolumeLoadObject: jest.fn(), + }, + imageLoader: { + createAndCacheDerivedImages: jest.fn(), + }, + volumeLoader: { + createAndCacheDerivedLabelmapVolume: jest.fn(), + createAndCacheVolumeFromImagesSync: jest.fn(), + }, + }; +}); + +jest.mock('../triggerSegmentationEvents', () => ({ + triggerSegmentationDataModified: jest.fn(), +})); + +import { moveSegmentToPrivateLabelmap } from './privateLabelmap'; +import { + beginLabelmapEditTransaction, + eraseLabelmapEditTransactionOverwrites, +} from './labelmapEditTransaction'; +import { getLabelmap } from './labelmapLayerStore'; +import { getSegmentBinding } from './labelmapSegmentBindings'; +import { createLabelmapMemo } from '../../../utilities/segmentation/createLabelmapMemo'; +import type { LabelmapMemo } from '../../../utilities/segmentation/createLabelmapMemo'; + +const { cache: mockCache, imageLoader: mockImageLoader } = jest.requireMock( + '@cornerstonejs/core' +); + +function createVoxelManager(values: number[]) { + const state = { values: [...values] }; + const voxelManager = { + id: `vm-${Math.random()}`, + dimensions: [values.length, 1, 1] as Types.Point3, + forEach( + callback: (args: { + value: number; + index: number; + pointIJK: Types.Point3; + }) => void + ) { + state.values.forEach((value, index) => { + callback({ value, index, pointIJK: [index, 0, 0] }); + }); + }, + getAtIndex: (index: number) => state.values[index], + setAtIndex: (index: number, value: number) => { + state.values[index] = value; + }, + getAtIJKPoint: ([i]: Types.Point3) => state.values[i], + setAtIJKPoint: ([i]: Types.Point3, value: number) => { + state.values[i] = value; + }, + get values() { + return [...state.values]; + }, + }; + return voxelManager as unknown as Types.IVoxelManager & { + values: number[]; + }; +} + +/** + * layerA holds segment 1 (label 1) and segment 2 (label 2); layerB holds + * segment 3 (label 1) - the post-overlap shape produced by a previous + * segment move. + */ +function createSegmentation(): Segmentation { + return { + segmentationId: 'segmentation', + label: 'Segmentation', + cachedStats: {}, + segments: { + 1: { + segmentIndex: 1, + label: 'S1', + locked: false, + cachedStats: {}, + active: false, + }, + 2: { + segmentIndex: 2, + label: 'S2', + locked: false, + cachedStats: {}, + active: true, + }, + 3: { + segmentIndex: 3, + label: 'S3', + locked: false, + cachedStats: {}, + active: false, + }, + }, + representationData: { + Labelmap: { + labelmaps: { + layerA: { + labelmapId: 'layerA', + storageKind: 'stack', + imageIds: ['layerA-image'], + labelToSegmentIndex: { 1: 1, 2: 2 }, + }, + layerB: { + labelmapId: 'layerB', + storageKind: 'stack', + imageIds: ['layerB-image'], + labelToSegmentIndex: { 1: 3 }, + }, + }, + segmentBindings: { + 1: { labelmapId: 'layerA', labelValue: 1 }, + 2: { labelmapId: 'layerA', labelValue: 2 }, + 3: { labelmapId: 'layerB', labelValue: 1 }, + }, + }, + }, + } as unknown as Segmentation; +} + +const imageData = { + indexToWorld: (point: Types.Point3) => point, +} as vtkImageData; + +/** Wires cache.getImage to serve the given imageId -> voxelManager map and + * makes derived-image creation produce the private layer's image. */ +function wireImages( + images: Record> +) { + mockCache.getImage.mockImplementation((imageId: string) => + images[imageId] ? { imageId, voxelManager: images[imageId] } : undefined + ); + mockImageLoader.createAndCacheDerivedImages.mockImplementation(() => { + const voxelManager = createVoxelManager( + new Array(Object.values(images)[0].values.length).fill(0) + ); + images['private-image'] = voxelManager; + return [{ imageId: 'private-image', voxelManager }]; + }); +} + +describe('overlap undo/redo bookkeeping', () => { + beforeEach(() => { + jest.clearAllMocks(); + }); + + describe('moveSegmentToPrivateLabelmap', () => { + it('moves the segment voxels and emits a step that reverses the whole move', () => { + const segmentation = createSegmentation(); + const layerAVoxels = createVoxelManager([1, 2, 2, 0]); + wireImages({ 'layerA-image': layerAVoxels }); + + let moveStep: LabelmapRestoreStep; + const privateLayer = moveSegmentToPrivateLabelmap(segmentation, 2, { + moveStepCallback: (step) => (moveStep = step), + }); + + const privateVoxels = () => + ( + mockCache.getImage('private-image')?.voxelManager as ReturnType< + typeof createVoxelManager + > + ).values; + + // the move itself: label 2 voxels leave layerA and land as label 1 on + // the private layer, and the binding follows + expect(privateLayer.labelmapId).not.toBe('layerA'); + expect(layerAVoxels.values).toEqual([1, 0, 0, 0]); + expect(privateVoxels()).toEqual([0, 1, 1, 0]); + expect(getSegmentBinding(segmentation, 2)).toEqual({ + labelmapId: privateLayer.labelmapId, + labelValue: 1, + }); + expect(moveStep).toBeDefined(); + + // undo: voxels return to layerA, binding restored, private layer gone + moveStep.undo(); + expect(layerAVoxels.values).toEqual([1, 2, 2, 0]); + expect(privateVoxels()).toEqual([0, 0, 0, 0]); + expect(getSegmentBinding(segmentation, 2)).toEqual({ + labelmapId: 'layerA', + labelValue: 2, + }); + expect( + getLabelmap(segmentation, privateLayer.labelmapId) + ).toBeUndefined(); + + // redo: the whole move replays, including re-registering the layer + moveStep.redo(); + expect(layerAVoxels.values).toEqual([1, 0, 0, 0]); + expect(privateVoxels()).toEqual([0, 1, 1, 0]); + expect(getSegmentBinding(segmentation, 2)).toEqual({ + labelmapId: privateLayer.labelmapId, + labelValue: 1, + }); + expect(getLabelmap(segmentation, privateLayer.labelmapId)).toBe( + privateLayer + ); + + // undo/redo stay stable across repeated cycles + moveStep.undo(); + expect(layerAVoxels.values).toEqual([1, 2, 2, 0]); + moveStep.redo(); + expect(privateVoxels()).toEqual([0, 1, 1, 0]); + }); + + it('does not emit a step when the segment is already alone on its layer', () => { + const segmentation = createSegmentation(); + const moveStepCallback = jest.fn(); + + const layer = moveSegmentToPrivateLabelmap(segmentation, 3, { + moveStepCallback, + }); + + expect(layer?.labelmapId).toBe('layerB'); + expect(moveStepCallback).not.toHaveBeenCalled(); + }); + }); + + describe('beginLabelmapEditTransaction', () => { + it('surfaces the move step when the default mover relocates the segment', () => { + const segmentation = createSegmentation(); + const layerAVoxels = createVoxelManager([1, 2, 0]); + wireImages({ 'layerA-image': layerAVoxels }); + + const transaction = beginLabelmapEditTransaction(segmentation, { + segmentIndex: 2, + overwriteSegmentIndices: [], + segmentationVoxelManager: layerAVoxels, + segmentationImageData: imageData, + isInObject: () => true, + }); + + expect(transaction.movedSegment).toBe(true); + expect(transaction.moveStep).toBeDefined(); + + // and the step round-trips + transaction.moveStep.undo(); + expect(layerAVoxels.values).toEqual([1, 2, 0]); + expect(getSegmentBinding(segmentation, 2)).toEqual({ + labelmapId: 'layerA', + labelValue: 2, + }); + + transaction.moveStep.redo(); + expect(layerAVoxels.values).toEqual([1, 0, 0]); + }); + + it('does not surface a move step when overlap is allowed (no move)', () => { + const segmentation = createSegmentation(); + const layerAVoxels = createVoxelManager([1, 2, 0]); + wireImages({ 'layerA-image': layerAVoxels }); + + const transaction = beginLabelmapEditTransaction(segmentation, { + segmentIndex: 2, + // overwrite of the sibling is allowed, so no protection, no move + overwriteSegmentIndices: [1], + segmentationVoxelManager: layerAVoxels, + segmentationImageData: imageData, + isInObject: () => true, + }); + + expect(transaction.movedSegment).toBe(false); + expect(transaction.moveStep).toBeUndefined(); + }); + }); + + describe('cross-layer erase recording', () => { + it('reports erased voxels through crossLayerEraseCallback so they can be restored', () => { + const segmentation = createSegmentation(); + const layerBVoxels = createVoxelManager([1, 1, 0]); + wireImages({ 'layerB-image': layerBVoxels }); + + const viewport = { + getCurrentImageId: () => 'ref-image', + getCurrentImageIdIndex: () => 0, + getImageIds: () => ['ref-image'], + } as unknown as Types.IStackViewport; + + const records = []; + const modifiedSlices = eraseLabelmapEditTransactionOverwrites( + segmentation, + { + segmentIndex: 2, + labelmapId: 'layerA', + labelValue: 2, + overwriteSegmentIndices: [3], + protectedSegmentIndices: [], + crossLayerEraseBindings: [{ labelmapId: 'layerB', labelValue: 1 }], + movedSegment: false, + }, + { + viewport, + referenceImageData: imageData, + isInObject: ([i]) => i < 2, + crossLayerEraseCallback: (record) => records.push(record), + } + ); + + expect(modifiedSlices).toEqual([0]); + expect(layerBVoxels.values).toEqual([0, 0, 0]); + expect(records).toHaveLength(1); + expect(records[0].labelValue).toBe(1); + expect(records[0].indices).toEqual([0, 1]); + + // undo semantics used by the strategy layer: write labelValue back + records[0].indices.forEach((index) => + records[0].voxelManager.setAtIndex(index, records[0].labelValue) + ); + expect(layerBVoxels.values).toEqual([1, 1, 0]); + }); + }); + + describe('full overlap stroke timeline', () => { + it('round-trips two strokes where the second moves the segment and erases another layer', () => { + const segmentation = createSegmentation(); + const layerAVoxels = createVoxelManager([1, 0, 0, 0]); + const layerBVoxels = createVoxelManager([0, 0, 1, 1]); + wireImages({ + 'layerA-image': layerAVoxels, + 'layerB-image': layerBVoxels, + }); + const viewport = { + getCurrentImageId: () => 'ref-image', + getCurrentImageIdIndex: () => 0, + getImageIds: () => ['ref-image'], + } as unknown as Types.IStackViewport; + + // --- stroke 1: segment 2 paints alone on the shared layer --- + const memoA = createLabelmapMemo( + 'segmentation', + layerAVoxels + ) as unknown as LabelmapMemo; + memoA.voxelManager.setAtIndex(1, 2); + expect(memoA.commitMemo()).toBe(true); + const afterStrokeA = [1, 2, 0, 0]; + expect(layerAVoxels.values).toEqual(afterStrokeA); + + // --- stroke 2: segment 2 paints over segment 1 with cross-layer + // overwrite of segment 3 - segment move + private write + erase --- + const transaction = beginLabelmapEditTransaction(segmentation, { + segmentIndex: 2, + overwriteSegmentIndices: [3], + segmentationVoxelManager: layerAVoxels, + segmentationImageData: imageData, + isInObject: () => true, + }); + expect(transaction.movedSegment).toBe(true); + + const privateLayerId = transaction.activeLayer.labelmapId; + const privateVM = mockCache.getImage('private-image') + .voxelManager as ReturnType; + // the segment move carried segment 2's stroke-1 voxel to the private layer + expect(privateVM.values).toEqual([0, 1, 0, 0]); + expect(layerAVoxels.values).toEqual([1, 0, 0, 0]); + + const memoB = createLabelmapMemo( + 'segmentation', + privateVM + ) as unknown as LabelmapMemo; + (memoB.priorSteps ||= []).push(transaction.moveStep); + + // the stroke's own writes land on the private layer + memoB.voxelManager.setAtIndex(0, 1); + + // cross-layer erase of segment 3, recorded as a post step - the same + // wiring crossLayerErase.ts performs + eraseLabelmapEditTransactionOverwrites(segmentation, transaction, { + viewport, + referenceImageData: imageData, + isInObject: ([i]) => i >= 2, + crossLayerEraseCallback: ({ voxelManager, labelValue, indices }) => { + (memoB.postSteps ||= []).push({ + undo: () => + indices.forEach((index) => + voxelManager.setAtIndex(index, labelValue) + ), + redo: () => + indices.forEach((index) => voxelManager.setAtIndex(index, 0)), + }); + }, + }); + expect(layerBVoxels.values).toEqual([0, 0, 0, 0]); + expect(memoB.commitMemo()).toBe(true); + + const finalPrivate = [1, 1, 0, 0]; + expect(privateVM.values).toEqual(finalPrivate); + + // --- undo stroke 2: private layer gone, segment 2 back on the shared + // layer with its stroke-1 voxel, segment 3 restored --- + memoB.restoreMemo(true); + expect(getLabelmap(segmentation, privateLayerId)).toBeUndefined(); + expect(getSegmentBinding(segmentation, 2)).toEqual({ + labelmapId: 'layerA', + labelValue: 2, + }); + expect(layerAVoxels.values).toEqual(afterStrokeA); + expect(layerBVoxels.values).toEqual([0, 0, 1, 1]); + expect(privateVM.values).toEqual([0, 0, 0, 0]); + + // --- undo stroke 1: back to the initial state --- + memoA.restoreMemo(true); + expect(layerAVoxels.values).toEqual([1, 0, 0, 0]); + + // --- redo both strokes: the exact final state returns --- + memoA.restoreMemo(false); + expect(layerAVoxels.values).toEqual(afterStrokeA); + + memoB.restoreMemo(false); + expect(getLabelmap(segmentation, privateLayerId)).toBe( + transaction.activeLayer + ); + expect(getSegmentBinding(segmentation, 2)).toEqual({ + labelmapId: privateLayerId, + labelValue: 1, + }); + expect(layerAVoxels.values).toEqual([1, 0, 0, 0]); + expect(privateVM.values).toEqual(finalPrivate); + expect(layerBVoxels.values).toEqual([0, 0, 0, 0]); + }); + }); +}); diff --git a/packages/tools/src/stateManagement/segmentation/labelmapModel/privateLabelmap.ts b/packages/tools/src/stateManagement/segmentation/labelmapModel/privateLabelmap.ts index 02f9f6d549..2f836efd1c 100644 --- a/packages/tools/src/stateManagement/segmentation/labelmapModel/privateLabelmap.ts +++ b/packages/tools/src/stateManagement/segmentation/labelmapModel/privateLabelmap.ts @@ -7,13 +7,34 @@ import { import type { Types } from '@cornerstonejs/core'; import type { Segmentation } from '../../../types/SegmentationStateTypes'; import type { LabelmapLayer } from '../../../types/LabelmapTypes'; -import { getLabelmap, registerLabelmap } from './labelmapLayerStore'; +import { + getLabelmap, + registerLabelmap, + removeLabelmap, +} from './labelmapLayerStore'; import { getSegmentBinding, getSegmentsOnLabelmap, setSegmentBinding, } from './labelmapSegmentBindings'; import { syncLegacyLabelmapData } from './labelmapLegacyAdapter'; +import type { LabelmapRestoreStep } from '../../../utilities/segmentation/createLabelmapMemo'; + +/** One source->target voxel-manager pair touched by a segment move, + * with the indices whose value moved. */ +type MovedVoxels = { + source: Types.IVoxelManager; + target: Types.IVoxelManager; + indices: number[]; +}; + +type MoveSegmentToPrivateLabelmapOptions = { + /** When provided, receives an undo/redo step that reverses/replays the whole + * move (layer registration, bulk voxel move, and binding change). + * The step does NOT fire segmentation events itself - the recorder is + * expected to wrap it with the appropriate data-modified notification. */ + moveStepCallback?: (step: LabelmapRestoreStep) => void; +}; function createPrivateVolumeLabelmap( segmentation: Segmentation, @@ -96,7 +117,8 @@ function createPrivateLabelmap( function moveSegmentToPrivateLabelmap( segmentation: Segmentation, - segmentIndex: number + segmentIndex: number, + options: MoveSegmentToPrivateLabelmapOptions = {} ): LabelmapLayer | undefined { const binding = getSegmentBinding(segmentation, segmentIndex); if (!binding) { @@ -117,9 +139,12 @@ function moveSegmentToPrivateLabelmap( const privateLabelmap = createPrivateLabelmap(segmentation, sourceLabelmap); registerLabelmap(segmentation, privateLabelmap); + const movedVoxels: MovedVoxels[] = []; + if (sourceLabelmap.volumeId && privateLabelmap.volumeId) { const sourceVolume = cache.getVolume(sourceLabelmap.volumeId); const targetVolume = cache.getVolume(privateLabelmap.volumeId); + const indices: number[] = []; sourceVolume.voxelManager.forEach(({ value, index }) => { if (value !== binding.labelValue) { return; @@ -127,7 +152,15 @@ function moveSegmentToPrivateLabelmap( targetVolume.voxelManager.setAtIndex(index, 1); sourceVolume.voxelManager.setAtIndex(index, 0); + indices.push(index); }); + if (indices.length) { + movedVoxels.push({ + source: sourceVolume.voxelManager as Types.IVoxelManager, + target: targetVolume.voxelManager as Types.IVoxelManager, + indices, + }); + } } else { const sourceImageIds = sourceLabelmap.imageIds ?? []; const targetImageIds = privateLabelmap.imageIds ?? []; @@ -140,6 +173,7 @@ function moveSegmentToPrivateLabelmap( return; } + const indices: number[] = []; sourceImage.voxelManager.forEach(({ value, index }) => { if (value !== binding.labelValue) { return; @@ -147,16 +181,56 @@ function moveSegmentToPrivateLabelmap( targetImage.voxelManager.setAtIndex(index, 1); sourceImage.voxelManager.setAtIndex(index, 0); + indices.push(index); }); + if (indices.length) { + movedVoxels.push({ + source: sourceImage.voxelManager as Types.IVoxelManager, + target: targetImage.voxelManager as Types.IVoxelManager, + indices, + }); + } }); } - setSegmentBinding(segmentation, segmentIndex, { + const previousBinding = { ...binding }; + const newBinding = { labelmapId: privateLabelmap.labelmapId, labelValue: 1, - }); + }; + + setSegmentBinding(segmentation, segmentIndex, { ...newBinding }); syncLegacyLabelmapData(segmentation); + // Hand the caller an undo/redo step for the WHOLE move so it can be + // recorded on the stroke's memo: without it, undo restores voxels but leaves + // the private layer, its binding, and the bulk-moved voxels behind, + // stranding the segment outside its history. + options.moveStepCallback?.({ + undo: () => { + for (const { source, target, indices } of movedVoxels) { + for (const index of indices) { + target.setAtIndex(index, 0); + source.setAtIndex(index, previousBinding.labelValue); + } + } + setSegmentBinding(segmentation, segmentIndex, { ...previousBinding }); + removeLabelmap(segmentation, privateLabelmap.labelmapId); + syncLegacyLabelmapData(segmentation); + }, + redo: () => { + registerLabelmap(segmentation, privateLabelmap); + for (const { source, target, indices } of movedVoxels) { + for (const index of indices) { + source.setAtIndex(index, 0); + target.setAtIndex(index, 1); + } + } + setSegmentBinding(segmentation, segmentIndex, { ...newBinding }); + syncLegacyLabelmapData(segmentation); + }, + }); + return privateLabelmap; } diff --git a/packages/tools/src/tools/segmentation/strategies/BrushStrategy.ts b/packages/tools/src/tools/segmentation/strategies/BrushStrategy.ts index a497ceae2f..457f6f3c94 100644 --- a/packages/tools/src/tools/segmentation/strategies/BrushStrategy.ts +++ b/packages/tools/src/tools/segmentation/strategies/BrushStrategy.ts @@ -7,6 +7,7 @@ import { StrategyCallbacks } from '../../../enums'; import type { LabelmapToolOperationDataAny } from '../../../types/LabelmapToolOperationData'; import type vtkImageData from '@kitware/vtk.js/Common/DataModel/ImageData'; import type { LabelmapMemo } from '../../../utilities/segmentation/createLabelmapMemo'; +import { memoAsStep } from '../../../utilities/segmentation/createLabelmapMemo'; import { eraseCrossLayerOverwrites, prepareOverlapOperationData, @@ -274,10 +275,41 @@ export default class BrushStrategy { initializedData.memo?.segmentationVoxelManager !== initializedData.segmentationVoxelManager ) { + // Mid-stroke voxel-manager swap (the segment moved to a private layer): + // commit the earlier same-stroke writes and ride them on the new memo so + // the whole stroke stays one undo/redo unit. + const previousMemo = initializedData.memo; initializedData.memo = initializedData.createMemo( initializedData.segmentationId, initializedData.segmentationVoxelManager ); + if ( + previousMemo && + previousMemo !== initializedData.memo && + previousMemo.commitMemo?.() + ) { + (initializedData.memo.priorSteps ||= []).push(memoAsStep(previousMemo)); + } + } + + // Record the segment move (layer registration + bulk voxel move + binding + // change) on the stroke's memo, so undo removes the private layer and + // returns the segment to its previous layer, and redo replays the move. + // The raw step doesn't fire events itself (keeps the labelmap model free + // of the event-layer import), so wrap it with the data-modified trigger. + const moveStep = initializedData.labelmapEditTransaction?.moveStep; + if (moveStep && initializedData.memo) { + const { segmentationId } = initializedData; + (initializedData.memo.priorSteps ||= []).push({ + undo: () => { + moveStep.undo(); + triggerSegmentationDataModified(segmentationId); + }, + redo: () => { + moveStep.redo(); + triggerSegmentationDataModified(segmentationId); + }, + }); } if ( diff --git a/packages/tools/src/tools/segmentation/strategies/utils/crossLayerErase.ts b/packages/tools/src/tools/segmentation/strategies/utils/crossLayerErase.ts index 92ccf70be4..917017c167 100644 --- a/packages/tools/src/tools/segmentation/strategies/utils/crossLayerErase.ts +++ b/packages/tools/src/tools/segmentation/strategies/utils/crossLayerErase.ts @@ -1,4 +1,5 @@ import { getSegmentation } from '../../../../stateManagement/segmentation/getSegmentation'; +import { triggerSegmentationDataModified } from '../../../../stateManagement/segmentation/triggerSegmentationEvents'; import { collectCrossLayerEraseBindings, eraseLabelmapEditTransactionOverwrites, @@ -29,6 +30,8 @@ function eraseCrossLayerOverwrites( return []; } + const { memo } = operationData; + return eraseLabelmapEditTransactionOverwrites( segmentation, operationData.labelmapEditTransaction, @@ -38,6 +41,27 @@ function eraseCrossLayerOverwrites( isInObject: operationData.isInObject, isInObjectBoundsIJK: operationData.isInObjectBoundsIJK, imageId: operationData.imageId, + // Record every cross-layer erase on the stroke's memo so undo/redo + // restores the other layers too - these writes bypass the memo's own + // history voxel manager (they target other layers' voxel managers). + crossLayerEraseCallback: memo + ? ({ voxelManager, labelValue, indices }) => { + (memo.postSteps ||= []).push({ + undo: () => { + for (const index of indices) { + voxelManager.setAtIndex(index, labelValue); + } + triggerSegmentationDataModified(operationData.segmentationId); + }, + redo: () => { + for (const index of indices) { + voxelManager.setAtIndex(index, 0); + } + triggerSegmentationDataModified(operationData.segmentationId); + }, + }); + } + : undefined, } ); } diff --git a/packages/tools/src/utilities/segmentation/createLabelmapMemo.spec.ts b/packages/tools/src/utilities/segmentation/createLabelmapMemo.spec.ts new file mode 100644 index 0000000000..80a0ac6790 --- /dev/null +++ b/packages/tools/src/utilities/segmentation/createLabelmapMemo.spec.ts @@ -0,0 +1,197 @@ +import type { Types } from '@cornerstonejs/core'; + +jest.mock( + '../../stateManagement/segmentation/triggerSegmentationEvents', + () => ({ + triggerSegmentationDataModified: jest.fn(), + }) +); + +import { createLabelmapMemo, memoAsStep } from './createLabelmapMemo'; +import type { LabelmapMemo, LabelmapRestoreStep } from './createLabelmapMemo'; +import { triggerSegmentationDataModified } from '../../stateManagement/segmentation/triggerSegmentationEvents'; + +/** + * A minimal segmentation voxel manager the RLE history voxel manager can wrap: + * dimensions [n, 1, 1] so index === pointIJK[0]. + */ +function createSourceVoxelManager(values: number[]) { + const state = { values: [...values] }; + const voxelManager = { + id: 'source', + dimensions: [values.length, 1, 1] as Types.Point3, + getAtIndex: (index: number) => state.values[index], + setAtIndex: (index: number, value: number) => { + state.values[index] = value; + }, + getAtIJKPoint: ([i]: Types.Point3) => state.values[i], + setAtIJKPoint: ([i]: Types.Point3, value: number) => { + state.values[i] = value; + }, + get values() { + return [...state.values]; + }, + }; + return voxelManager as unknown as Types.IVoxelManager & { + values: number[]; + }; +} + +function createOrderedStep( + label: string, + order: string[] +): LabelmapRestoreStep { + return { + undo: () => order.push(`${label}.undo`), + redo: () => order.push(`${label}.redo`), + }; +} + +describe('createLabelmapMemo', () => { + beforeEach(() => { + jest.clearAllMocks(); + }); + + describe('without overlap (plain single-layer stroke)', () => { + it('round-trips a stroke through undo and redo', () => { + const source = createSourceVoxelManager([0, 0, 2, 0]); + const memo = createLabelmapMemo( + 'segmentation', + source + ) as unknown as LabelmapMemo; + + // The stroke writes through the memo's history voxel manager, exactly + // like the brush strategies do. + memo.voxelManager.setAtIndex(0, 1); + memo.voxelManager.setAtIndex(2, 1); + expect(source.values).toEqual([1, 0, 1, 0]); + + expect(memo.commitMemo()).toBe(true); + + memo.restoreMemo(true); + expect(source.values).toEqual([0, 0, 2, 0]); + + memo.restoreMemo(false); + expect(source.values).toEqual([1, 0, 1, 0]); + + // undo/redo must stay stable across repeated cycles + memo.restoreMemo(true); + expect(source.values).toEqual([0, 0, 2, 0]); + memo.restoreMemo(false); + expect(source.values).toEqual([1, 0, 1, 0]); + + expect(triggerSegmentationDataModified).toHaveBeenCalledWith( + 'segmentation', + expect.any(Array) + ); + }); + + it('does not commit a stroke that changed nothing', () => { + const source = createSourceVoxelManager([0, 1, 0]); + const memo = createLabelmapMemo( + 'segmentation', + source + ) as unknown as LabelmapMemo; + + // writing the value that is already there is a no-op + memo.voxelManager.setAtIndex(1, 1); + + expect(memo.commitMemo()).toBe(false); + }); + + it('stays committed once committed', () => { + const source = createSourceVoxelManager([0]); + const memo = createLabelmapMemo( + 'segmentation', + source + ) as unknown as LabelmapMemo; + + memo.voxelManager.setAtIndex(0, 1); + expect(memo.commitMemo()).toBe(true); + expect(memo.commitMemo()).toBe(true); + }); + }); + + describe('with overlap (steps riding on the stroke memo)', () => { + it('commits a memo that only recorded steps', () => { + const source = createSourceVoxelManager([0]); + const memo = createLabelmapMemo( + 'segmentation', + source + ) as unknown as LabelmapMemo; + const order: string[] = []; + + memo.priorSteps = [createOrderedStep('segmentMove', order)]; + + expect(memo.commitMemo()).toBe(true); + + // A steps-only memo must restore without its own undo/redo voxel + // managers. + memo.restoreMemo(true); + expect(order).toEqual(['segmentMove.undo']); + memo.restoreMemo(false); + expect(order).toEqual(['segmentMove.undo', 'segmentMove.redo']); + }); + + it('restores steps in reverse-chronological order on undo and chronological order on redo', () => { + const source = createSourceVoxelManager([0, 0]); + const memo = createLabelmapMemo( + 'segmentation', + source + ) as unknown as LabelmapMemo; + const order: string[] = []; + + // one voxel write so the own-voxel restore is a single, observable call + memo.voxelManager.setAtIndex(0, 1); + + memo.priorSteps = [ + createOrderedStep('earlierMemo', order), + createOrderedStep('segmentMove', order), + ]; + memo.postSteps = [createOrderedStep('crossLayerErase', order)]; + + expect(memo.commitMemo()).toBe(true); + + const originalSetAtIJKPoint = source.setAtIJKPoint; + source.setAtIJKPoint = (point, value) => { + order.push('ownVoxels'); + return originalSetAtIJKPoint(point, value); + }; + + memo.restoreMemo(true); + expect(order).toEqual([ + 'crossLayerErase.undo', + 'ownVoxels', + 'segmentMove.undo', + 'earlierMemo.undo', + ]); + + order.length = 0; + memo.restoreMemo(false); + expect(order).toEqual([ + 'earlierMemo.redo', + 'segmentMove.redo', + 'ownVoxels', + 'crossLayerErase.redo', + ]); + }); + + it('memoAsStep replays the wrapped memo', () => { + const source = createSourceVoxelManager([0]); + const memo = createLabelmapMemo( + 'segmentation', + source + ) as unknown as LabelmapMemo; + + memo.voxelManager.setAtIndex(0, 3); + expect(memo.commitMemo()).toBe(true); + + const step = memoAsStep(memo); + + step.undo(); + expect(source.values).toEqual([0]); + step.redo(); + expect(source.values).toEqual([3]); + }); + }); +}); diff --git a/packages/tools/src/utilities/segmentation/createLabelmapMemo.ts b/packages/tools/src/utilities/segmentation/createLabelmapMemo.ts index 3c4ab32cdc..28989572fe 100644 --- a/packages/tools/src/utilities/segmentation/createLabelmapMemo.ts +++ b/packages/tools/src/utilities/segmentation/createLabelmapMemo.ts @@ -4,6 +4,18 @@ import type { Types } from '@cornerstonejs/core'; const { VoxelManager, RLEVoxelMap } = utilities; +/** + * A side effect of a stroke that is not captured by the memo's own history + * voxel manager: cross-layer erases, a segment moving to a private + * labelmap layer (layer registration + binding change + bulk voxel move), or a + * committed sibling memo from before a mid-stroke voxel-manager swap. Steps + * make the whole stroke restorable as one undo/redo unit. + */ +export type LabelmapRestoreStep = { + undo: () => void; + redo: () => void; +}; + /** * The labelmap memo state, extending from the base Memo state */ @@ -18,6 +30,17 @@ export type LabelmapMemo = Types.Memo & { memo?: LabelmapMemo; /** A unique identifier for this memo */ id: string; + /** + * Steps that happened chronologically BEFORE this memo's voxel writes + * (earlier same-stroke memo on another layer, a segment move to a + * private layer). + */ + priorSteps?: LabelmapRestoreStep[]; + /** + * Steps that happened chronologically AFTER/DURING this memo's voxel writes + * (cross-layer overwrite erases). + */ + postSteps?: LabelmapRestoreStep[]; }; /** @@ -32,19 +55,60 @@ export function createLabelmapMemo( } /** - * A restore memo function. This simply copies either the redo or the base - * voxel manager data to the segmentation state and triggers segmentation data - * modified. + * Wraps a committed memo so it can ride along as a step of another memo (used + * when a stroke's target voxel manager swaps mid-stroke, e.g. a segment moves + * to a private labelmap layer part-way through a drag). + */ +export function memoAsStep(memo: LabelmapMemo): LabelmapRestoreStep { + return { + undo: () => memo.restoreMemo(true), + redo: () => memo.restoreMemo(false), + }; +} + +/** + * A restore memo function. This restores the memo's own voxel changes plus + * any recorded steps, in reverse-chronological order for undo and + * chronological order for redo, so a stroke that moved a segment across + * layers or erased other layers restores as one consistent unit. */ export function restoreMemo(isUndo?: boolean) { - const { segmentationVoxelManager, undoVoxelManager, redoVoxelManager } = this; - const useVoxelManager = - isUndo === false ? redoVoxelManager : undoVoxelManager; - useVoxelManager.forEach(({ value, pointIJK }) => { - segmentationVoxelManager.setAtIJKPoint(pointIJK, value); - }); - const slices = useVoxelManager.getArrayOfModifiedSlices(); - triggerSegmentationDataModified(this.segmentationId, slices); + const undo = isUndo !== false; + const priorSteps: LabelmapRestoreStep[] = this.priorSteps ?? []; + const postSteps: LabelmapRestoreStep[] = this.postSteps ?? []; + + const restoreVoxelChanges = () => { + const { segmentationVoxelManager, undoVoxelManager, redoVoxelManager } = + this; + const useVoxelManager = undo ? undoVoxelManager : redoVoxelManager; + if (!useVoxelManager) { + // Steps-only memo (no voxel writes of its own) + return; + } + useVoxelManager.forEach(({ value, pointIJK }) => { + segmentationVoxelManager.setAtIJKPoint(pointIJK, value); + }); + const slices = useVoxelManager.getArrayOfModifiedSlices(); + triggerSegmentationDataModified(this.segmentationId, slices); + }; + + if (undo) { + for (let i = postSteps.length - 1; i >= 0; i--) { + postSteps[i].undo(); + } + restoreVoxelChanges(); + for (let i = priorSteps.length - 1; i >= 0; i--) { + priorSteps[i].undo(); + } + } else { + for (const step of priorSteps) { + step.redo(); + } + restoreVoxelChanges(); + for (const step of postSteps) { + step.redo(); + } + } // Event dispatch moved to historyMemo/index.ts } @@ -80,8 +144,11 @@ function commitMemo() { if (this.redoVoxelManager) { return true; } + const hasSteps = !!(this.priorSteps?.length || this.postSteps?.length); if (!this.voxelManager.modifiedSlices.size) { - return false; + // No voxel writes of its own, but recorded steps still make this memo + // worth keeping on the history ring. + return hasSteps; } const { segmentationVoxelManager } = this; const undoVoxelManager = VoxelManager.createRLEHistoryVoxelManager( From 44783f7b7797fedd425a52c08d1a033a87cbd5fb Mon Sep 17 00:00:00 2001 From: Alireza Date: Mon, 6 Jul 2026 11:38:01 -0400 Subject: [PATCH 2/5] perf(segmentation): memoize labelmap image reference scans Every eager (non-lazy) brush or eraser drag event re-ran isReferenceViewable over ALL labelmap imageIds via updateLabelmapSegmentationImageReferences, each call doing full metadata and slice-basis resolution. On a two-layer, 120-image labelmap this cost about 51ms per drag event (about 19fps while erasing). The resolver now remembers completed scans per segmentation, keyed by the viewport's current image plus a fingerprint of the labelmap imageId set, and replays the stored mapping instead of rescanning. Layer additions or removals change the fingerprint and force a fresh scan; removeSegmentation and reset clear the memo. Same scenario now costs about 2ms per event. --- .../labelmapImageReferenceResolver.ts | 80 ++++++++++++++++++- 1 file changed, 78 insertions(+), 2 deletions(-) diff --git a/packages/tools/src/stateManagement/segmentation/labelmapModel/labelmapImageReferenceResolver.ts b/packages/tools/src/stateManagement/segmentation/labelmapModel/labelmapImageReferenceResolver.ts index 2e52698310..afd3489148 100644 --- a/packages/tools/src/stateManagement/segmentation/labelmapModel/labelmapImageReferenceResolver.ts +++ b/packages/tools/src/stateManagement/segmentation/labelmapModel/labelmapImageReferenceResolver.ts @@ -31,6 +31,19 @@ class LabelmapImageReferenceResolver { // labelmapImageIdReferenceMap, so removeSegmentation can purge entries in // O(k) without scanning the whole map. private readonly keysBySegmentationId = new Map>(); + // Memo of already-performed isReferenceViewable scans. Each scan resolves + // every labelmap imageId against the viewport's current image via full + // metadata resolution (image plane, slice basis), which costs tens of ms + // for multi-layer labelmaps - and the segmentation data-modified pipeline + // re-enters this on EVERY brush/erase drag event. The result only depends + // on (current image, labelmap imageId set), so scans are replayed from + // stackLabelmapImageIdReferenceMap once their key has been seen. Keys embed + // a signature of the labelmap imageId set, so adding/removing a layer + // (which changes the set) naturally forces a fresh scan. + private readonly completedScanKeysBySegmentationId = new Map< + string, + Set + >(); constructor(getSegmentation: GetSegmentation) { this.getSegmentation = getSegmentation; @@ -40,10 +53,12 @@ class LabelmapImageReferenceResolver { this.stackLabelmapImageIdReferenceMap.clear(); this.labelmapImageIdReferenceMap.clear(); this.keysBySegmentationId.clear(); + this.completedScanKeysBySegmentationId.clear(); } removeSegmentation(segmentationId: string): void { this.stackLabelmapImageIdReferenceMap.delete(segmentationId); + this.completedScanKeysBySegmentationId.delete(segmentationId); const keys = this.keysBySegmentationId.get(segmentationId); if (keys) { for (const key of keys) { @@ -142,11 +157,60 @@ class LabelmapImageReferenceResolver { return; } - return this.updateLabelmapSegmentationReferences( + const viewport = enabledElement.viewport as Types.IStackViewport; + const scanKey = this.generateScanKey('current', viewport, labelmapImageIds); + if (scanKey && this.hasCompletedScan(segmentationId, scanKey)) { + return this.stackLabelmapImageIdReferenceMap + .get(segmentationId) + ?.get(viewport.getCurrentImageId()); + } + + const result = this.updateLabelmapSegmentationReferences( segmentationId, - enabledElement.viewport as Types.IStackViewport, + viewport, labelmapImageIds ); + + if (scanKey) { + this.markCompletedScan(segmentationId, scanKey); + } + + return result; + } + + /** Signature of one isReferenceViewable scan: the viewport's current image + * plus a cheap fingerprint of the labelmap imageId set (count + endpoints), + * so layer additions/removals invalidate naturally. */ + private generateScanKey( + kind: 'current' | 'all', + viewport: Types.IStackViewport, + labelmapImageIds: string[] + ): string | undefined { + const referenceImageId = + typeof viewport.getCurrentImageId === 'function' + ? viewport.getCurrentImageId() + : undefined; + if (!referenceImageId) { + return; + } + return `${kind}|${viewport.id}|${referenceImageId}|${ + labelmapImageIds.length + }|${labelmapImageIds[0]}|${labelmapImageIds[labelmapImageIds.length - 1]}`; + } + + private hasCompletedScan(segmentationId: string, scanKey: string): boolean { + return !!this.completedScanKeysBySegmentationId + .get(segmentationId) + ?.has(scanKey); + } + + private markCompletedScan(segmentationId: string, scanKey: string): void { + let keys = this.completedScanKeysBySegmentationId.get(segmentationId); + if (!keys) { + keys = new Set(); + this.completedScanKeysBySegmentationId.set(segmentationId, keys); + } + keys.add(scanKey); } getCurrentLabelmapImageIdsForViewport( @@ -382,6 +446,18 @@ class LabelmapImageReferenceResolver { const stackViewport = enabledElement.viewport as Types.IStackViewport; + const scanKey = this.generateScanKey( + 'all', + stackViewport, + labelmapImageIds + ); + if (scanKey && this.hasCompletedScan(segmentationId, scanKey)) { + return; + } + if (scanKey) { + this.markCompletedScan(segmentationId, scanKey); + } + this.updateLabelmapSegmentationReferences( segmentationId, stackViewport, From 3489cbfe2dccaa785de2cfe5b83f14680f023070 Mon Sep 17 00:00:00 2001 From: Alireza Date: Mon, 6 Jul 2026 11:38:11 -0400 Subject: [PATCH 3/5] fix(BrushTool): bail out of preMouseDown when hover data cannot be created preMouseDownCallback dereferenced this._hoverData.viewport without guarding the case where createHoverData returns undefined (no active segment index yet, or the viewport camera is not resolvable), crashing with 'Cannot read properties of undefined (reading viewport)'. Every other createHoverData call site already guards this; do the same here and undo the draw activation before bailing. --- packages/tools/src/tools/segmentation/BrushTool.ts | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/packages/tools/src/tools/segmentation/BrushTool.ts b/packages/tools/src/tools/segmentation/BrushTool.ts index bc8d8452f3..2376a61fa4 100644 --- a/packages/tools/src/tools/segmentation/BrushTool.ts +++ b/packages/tools/src/tools/segmentation/BrushTool.ts @@ -281,6 +281,15 @@ class BrushTool extends LabelmapBaseTool { }; this._hoverData = this.createHoverData(element, canvasPoint); + if (!this._hoverData) { + // Transient state (e.g. no active segment index yet, or the viewport's + // camera is not resolvable) - undo the draw activation done above and + // bail like the other createHoverData call sites do. + this._deactivateDraw(element); + resetElementCursor(element); + this._editData = null; + return false; + } this._calculateCursor(element, canvasPoint); this._resetLazyEditState(); From 7f82dd8d179ba20e0b0ac998d2c2a9a38251d679 Mon Sep 17 00:00:00 2001 From: Alireza Date: Mon, 6 Jul 2026 11:59:03 -0400 Subject: [PATCH 4/5] fix(tools): cache labelmap scan results by viewport --- .../labelmapImageReferenceResolver.spec.ts | 34 +++++++++++++ .../labelmapImageReferenceResolver.ts | 49 +++++++++++-------- 2 files changed, 63 insertions(+), 20 deletions(-) diff --git a/packages/tools/src/stateManagement/segmentation/labelmapModel/labelmapImageReferenceResolver.spec.ts b/packages/tools/src/stateManagement/segmentation/labelmapModel/labelmapImageReferenceResolver.spec.ts index 66bc0b85d0..62ee3d3380 100644 --- a/packages/tools/src/stateManagement/segmentation/labelmapModel/labelmapImageReferenceResolver.spec.ts +++ b/packages/tools/src/stateManagement/segmentation/labelmapModel/labelmapImageReferenceResolver.spec.ts @@ -215,6 +215,40 @@ describe('LabelmapImageReferenceResolver', () => { ); expect(result).toBe('lm-a-0'); }); + + it('replays the cached result for the matching viewport scan key', () => { + const segmentation = makeSegmentation(); + getSegmentation.mockReturnValue(segmentation); + + const viewport1 = makeStackViewport({ + id: 'vp-1', + getCurrentImageId: jest.fn(() => 'ref-0'), + isReferenceViewable: jest.fn( + (ref) => ref.referencedImageId === 'lm-a-0' + ), + }); + const viewport2 = makeStackViewport({ + id: 'vp-2', + getCurrentImageId: jest.fn(() => 'ref-0'), + isReferenceViewable: jest.fn( + (ref) => ref.referencedImageId === 'lm-a-1' + ), + }); + getEnabledElementByViewportId.mockImplementation((viewportId) => ({ + viewport: viewportId === 'vp-1' ? viewport1 : viewport2, + })); + + expect( + resolver.updateLabelmapSegmentationImageReferences('vp-1', 'seg-1') + ).toBe('lm-a-0'); + expect( + resolver.updateLabelmapSegmentationImageReferences('vp-2', 'seg-1') + ).toBe('lm-a-1'); + expect( + resolver.updateLabelmapSegmentationImageReferences('vp-1', 'seg-1') + ).toBe('lm-a-0'); + expect(viewport1.isReferenceViewable).toHaveBeenCalledTimes(2); + }); }); describe('getCurrentLabelmapImageIdsForViewport', () => { diff --git a/packages/tools/src/stateManagement/segmentation/labelmapModel/labelmapImageReferenceResolver.ts b/packages/tools/src/stateManagement/segmentation/labelmapModel/labelmapImageReferenceResolver.ts index afd3489148..ae4e62c76b 100644 --- a/packages/tools/src/stateManagement/segmentation/labelmapModel/labelmapImageReferenceResolver.ts +++ b/packages/tools/src/stateManagement/segmentation/labelmapModel/labelmapImageReferenceResolver.ts @@ -35,14 +35,12 @@ class LabelmapImageReferenceResolver { // every labelmap imageId against the viewport's current image via full // metadata resolution (image plane, slice basis), which costs tens of ms // for multi-layer labelmaps - and the segmentation data-modified pipeline - // re-enters this on EVERY brush/erase drag event. The result only depends - // on (current image, labelmap imageId set), so scans are replayed from - // stackLabelmapImageIdReferenceMap once their key has been seen. Keys embed - // a signature of the labelmap imageId set, so adding/removing a layer - // (which changes the set) naturally forces a fresh scan. - private readonly completedScanKeysBySegmentationId = new Map< + // re-enters this on EVERY brush/erase drag event. Scan keys include the + // viewport id and a signature of the labelmap imageId set, so adding/removing + // a layer (which changes the set) naturally forces a fresh scan. + private readonly completedScanResultsBySegmentationId = new Map< string, - Set + Map >(); constructor(getSegmentation: GetSegmentation) { @@ -53,12 +51,12 @@ class LabelmapImageReferenceResolver { this.stackLabelmapImageIdReferenceMap.clear(); this.labelmapImageIdReferenceMap.clear(); this.keysBySegmentationId.clear(); - this.completedScanKeysBySegmentationId.clear(); + this.completedScanResultsBySegmentationId.clear(); } removeSegmentation(segmentationId: string): void { this.stackLabelmapImageIdReferenceMap.delete(segmentationId); - this.completedScanKeysBySegmentationId.delete(segmentationId); + this.completedScanResultsBySegmentationId.delete(segmentationId); const keys = this.keysBySegmentationId.get(segmentationId); if (keys) { for (const key of keys) { @@ -160,9 +158,7 @@ class LabelmapImageReferenceResolver { const viewport = enabledElement.viewport as Types.IStackViewport; const scanKey = this.generateScanKey('current', viewport, labelmapImageIds); if (scanKey && this.hasCompletedScan(segmentationId, scanKey)) { - return this.stackLabelmapImageIdReferenceMap - .get(segmentationId) - ?.get(viewport.getCurrentImageId()); + return this.getCompletedScanResult(segmentationId, scanKey); } const result = this.updateLabelmapSegmentationReferences( @@ -172,7 +168,7 @@ class LabelmapImageReferenceResolver { ); if (scanKey) { - this.markCompletedScan(segmentationId, scanKey); + this.markCompletedScan(segmentationId, scanKey, result); } return result; @@ -199,18 +195,31 @@ class LabelmapImageReferenceResolver { } private hasCompletedScan(segmentationId: string, scanKey: string): boolean { - return !!this.completedScanKeysBySegmentationId + return !!this.completedScanResultsBySegmentationId .get(segmentationId) ?.has(scanKey); } - private markCompletedScan(segmentationId: string, scanKey: string): void { - let keys = this.completedScanKeysBySegmentationId.get(segmentationId); - if (!keys) { - keys = new Set(); - this.completedScanKeysBySegmentationId.set(segmentationId, keys); + private getCompletedScanResult( + segmentationId: string, + scanKey: string + ): string | undefined { + return this.completedScanResultsBySegmentationId + .get(segmentationId) + ?.get(scanKey); + } + + private markCompletedScan( + segmentationId: string, + scanKey: string, + result?: string + ): void { + let results = this.completedScanResultsBySegmentationId.get(segmentationId); + if (!results) { + results = new Map(); + this.completedScanResultsBySegmentationId.set(segmentationId, results); } - keys.add(scanKey); + results.set(scanKey, result); } getCurrentLabelmapImageIdsForViewport( From cbdea49ba62b2c2d5b9cca9ed08278effcba4c17 Mon Sep 17 00:00:00 2001 From: Alireza Date: Mon, 6 Jul 2026 12:06:15 -0400 Subject: [PATCH 5/5] test(tools): assert viewport scan replay coverage --- .../labelmapModel/labelmapImageReferenceResolver.spec.ts | 1 + 1 file changed, 1 insertion(+) diff --git a/packages/tools/src/stateManagement/segmentation/labelmapModel/labelmapImageReferenceResolver.spec.ts b/packages/tools/src/stateManagement/segmentation/labelmapModel/labelmapImageReferenceResolver.spec.ts index 62ee3d3380..077f9aad22 100644 --- a/packages/tools/src/stateManagement/segmentation/labelmapModel/labelmapImageReferenceResolver.spec.ts +++ b/packages/tools/src/stateManagement/segmentation/labelmapModel/labelmapImageReferenceResolver.spec.ts @@ -248,6 +248,7 @@ describe('LabelmapImageReferenceResolver', () => { resolver.updateLabelmapSegmentationImageReferences('vp-1', 'seg-1') ).toBe('lm-a-0'); expect(viewport1.isReferenceViewable).toHaveBeenCalledTimes(2); + expect(viewport2.isReferenceViewable).toHaveBeenCalledTimes(2); }); });