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/labelmapImageReferenceResolver.spec.ts b/packages/tools/src/stateManagement/segmentation/labelmapModel/labelmapImageReferenceResolver.spec.ts index 66bc0b85d0..077f9aad22 100644 --- a/packages/tools/src/stateManagement/segmentation/labelmapModel/labelmapImageReferenceResolver.spec.ts +++ b/packages/tools/src/stateManagement/segmentation/labelmapModel/labelmapImageReferenceResolver.spec.ts @@ -215,6 +215,41 @@ 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); + expect(viewport2.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 2e52698310..ae4e62c76b 100644 --- a/packages/tools/src/stateManagement/segmentation/labelmapModel/labelmapImageReferenceResolver.ts +++ b/packages/tools/src/stateManagement/segmentation/labelmapModel/labelmapImageReferenceResolver.ts @@ -31,6 +31,17 @@ 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. 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, + Map + >(); constructor(getSegmentation: GetSegmentation) { this.getSegmentation = getSegmentation; @@ -40,10 +51,12 @@ class LabelmapImageReferenceResolver { this.stackLabelmapImageIdReferenceMap.clear(); this.labelmapImageIdReferenceMap.clear(); this.keysBySegmentationId.clear(); + this.completedScanResultsBySegmentationId.clear(); } removeSegmentation(segmentationId: string): void { this.stackLabelmapImageIdReferenceMap.delete(segmentationId); + this.completedScanResultsBySegmentationId.delete(segmentationId); const keys = this.keysBySegmentationId.get(segmentationId); if (keys) { for (const key of keys) { @@ -142,11 +155,71 @@ 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.getCompletedScanResult(segmentationId, scanKey); + } + + const result = this.updateLabelmapSegmentationReferences( segmentationId, - enabledElement.viewport as Types.IStackViewport, + viewport, labelmapImageIds ); + + if (scanKey) { + this.markCompletedScan(segmentationId, scanKey, result); + } + + 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.completedScanResultsBySegmentationId + .get(segmentationId) + ?.has(scanKey); + } + + 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); + } + results.set(scanKey, result); } getCurrentLabelmapImageIdsForViewport( @@ -382,6 +455,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, 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/BrushTool.ts b/packages/tools/src/tools/segmentation/BrushTool.ts index 9bcacebf09..6bcdde7a3b 100644 --- a/packages/tools/src/tools/segmentation/BrushTool.ts +++ b/packages/tools/src/tools/segmentation/BrushTool.ts @@ -277,6 +277,10 @@ 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) - clear the edit data assigned above and bail + // before the draw listeners are bound below. + this._editData = null; return false; } this._calculateCursor(element, canvasPoint); 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(