Fix volume undo for unfetched data

frontend/javascripts/oxalis/model/actions/volumetracing_actions.js

@@ -21,10 +21,12 @@ export type CopySegmentationLayerAction = {
   type: "COPY_SEGMENTATION_LAYER",
   source: "previousLayer" | "nextLayer",
 };
+export type MaybeBucketLoadedPromise = null | Promise<BucketDataArray>;
 export type AddBucketToUndoAction = {
   type: "ADD_BUCKET_TO_UNDO",
   zoomedBucketAddress: Vector4,
   bucketData: BucketDataArray,
+  maybeBucketLoadedPromise: MaybeBucketLoadedPromise,
 };
 type UpdateDirectionAction = { type: "UPDATE_DIRECTION", centroid: Vector3 };
 type ResetContourAction = { type: "RESET_CONTOUR" };
@@ -141,7 +143,13 @@ export const setContourTracingModeAction = (mode: ContourMode): SetContourTracin
 export const addBucketToUndoAction = (
   zoomedBucketAddress: Vector4,
   bucketData: BucketDataArray,
-): AddBucketToUndoAction => ({ type: "ADD_BUCKET_TO_UNDO", zoomedBucketAddress, bucketData });
+  maybeBucketLoadedPromise: MaybeBucketLoadedPromise,
+): AddBucketToUndoAction => ({
+  type: "ADD_BUCKET_TO_UNDO",
+  zoomedBucketAddress,
+  bucketData,
+  maybeBucketLoadedPromise,
+});
 
 export const inferSegmentationInViewportAction = (
   position: Vector3,

frontend/javascripts/oxalis/model/bucket_data_handling/bucket.js

@@ -240,15 +240,15 @@ export class DataBucket {
     return { isVoxelOutside, neighbourBucketAddress, adjustedVoxel };
   };
 
-  getCopyOfData(): BucketDataArray {
-    const bucketData = this.getOrCreateData();
+  getCopyOfData(): { dataClone: BucketDataArray, triggeredBucketFetch: boolean } {
+    const { data: bucketData, triggeredBucketFetch } = this.getOrCreateData();
     const TypedArrayClass = getConstructorForElementClass(this.elementClass)[0];
     const dataClone = new TypedArrayClass(bucketData);
-    return dataClone;
+    return { dataClone, triggeredBucketFetch };
   }
 
   label(labelFunc: BucketDataArray => void) {
-    const bucketData = this.getOrCreateData();
+    const { data: bucketData } = this.getOrCreateData();
     this.markAndAddBucketForUndo();
     labelFunc(bucketData);
     this.throttledTriggerLabeled();
@@ -258,7 +258,16 @@ export class DataBucket {
     this.dirty = true;
     if (!bucketsAlreadyInUndoState.has(this)) {
       bucketsAlreadyInUndoState.add(this);
-      Store.dispatch(addBucketToUndoAction(this.zoomedAddress, this.getCopyOfData()));
+      let maybeBucketLoadedPromise = null;
+      const { dataClone, triggeredBucketFetch } = this.getCopyOfData();
+      if (triggeredBucketFetch) {
+        maybeBucketLoadedPromise = new Promise((resolve, _reject) => {
+          this.once("bucketLoaded", backendData => resolve(backendData));
+        });
+      }
+      Store.dispatch(
+        addBucketToUndoAction(this.zoomedAddress, dataClone, maybeBucketLoadedPromise),
+      );
     }
   }
 
@@ -294,16 +303,17 @@ export class DataBucket {
     this.accessed = false;
   }
 
-  getOrCreateData(): BucketDataArray {
+  getOrCreateData(): { data: BucketDataArray, triggeredBucketFetch: boolean } {
+    let triggeredBucketFetch = false;
     if (this.data == null) {
       const [TypedArrayClass, channelCount] = getConstructorForElementClass(this.elementClass);
       this.data = new TypedArrayClass(channelCount * Constants.BUCKET_SIZE);
       if (!this.isMissing()) {
+        triggeredBucketFetch = true;
         this.temporalBucketManager.addBucket(this);
       }
     }
-
-    return this.getData();
+    return { data: this.getData(), triggeredBucketFetch };
   }
 
   pull(): void {
@@ -346,7 +356,7 @@ export class DataBucket {
         } else {
           this.data = data;
         }
-        this.trigger("bucketLoaded");
+        this.trigger("bucketLoaded", data);
         this.state = BucketStateEnum.LOADED;
         break;
       default:

frontend/javascripts/oxalis/model/bucket_data_handling/data_cube.js

@@ -412,7 +412,7 @@ class DataCube {
       );
     }
     const seedVoxelIndex = this.getVoxelIndex(seedVoxel, zoomStep);
-    const sourceCellId = seedBucket.getOrCreateData()[seedVoxelIndex];
+    const sourceCellId = seedBucket.getOrCreateData().data[seedVoxelIndex];
     if (sourceCellId === cellId) {
       return null;
     }
@@ -428,7 +428,7 @@ class DataCube {
       ) {
         continue;
       }
-      const bucketData = currentBucket.getOrCreateData();
+      const { data: bucketData } = currentBucket.getOrCreateData();
       const initialVoxelIndex = this.getVoxelIndexByVoxelOffset(initialVoxelInBucket);
       if (bucketData[initialVoxelIndex] !== sourceCellId) {
         // Ignoring neighbour buckets whose cellId at the initial voxel does not match the source cell id.

frontend/javascripts/oxalis/model/sagas/save_saga.js

@@ -21,6 +21,7 @@ import {
   type AddBucketToUndoAction,
   type FinishAnnotationStrokeAction,
   type ImportVolumeTracingAction,
+  type MaybeBucketLoadedPromise,
 } from "oxalis/model/actions/volumetracing_actions";
 import {
   _all,
@@ -75,7 +76,12 @@ import { enforceSkeletonTracing } from "../accessors/skeletontracing_accessor";
 
 const byteArrayToLz4Array = createWorker(compressLz4Block);
 
-type UndoBucket = { zoomedBucketAddress: Vector4, data: Uint8Array };
+type UndoBucket = {
+  zoomedBucketAddress: Vector4,
+  data: Uint8Array,
+  backendData?: Uint8Array,
+  maybeBucketLoadedPromise: MaybeBucketLoadedPromise,
+};
 type VolumeAnnotationBatch = Array<UndoBucket>;
 type SkeletonUndoState = { type: "skeleton", data: SkeletonTracing };
 type VolumeUndoState = { type: "volume", data: VolumeAnnotationBatch };
@@ -91,6 +97,14 @@ type racedActionsNeededForUndoRedo = {
   redo: ?RedoAction,
 };
 
+// Better version: Let the undo stack / batch know that bucket data is not yet fetched by the backend.
+// Then if the data is fetched backend, a promise passed to the saga as an argument of the ADD_BUCKET_TO_UNDO action
+// is triggered, that contains a reference to the compressed ground truth data of the backend and is then added to the bucket-data
+// in the undo stack.
+// Later when an undo action is triggered and a volume action should be undone, it is first tested whether this undo bucket has
+// ground truth data, and if yes, both are decompressed and then merged and applied.
+// If there is no ground truth, then proceed just like it is done currently.
+
 export function* collectUndoStates(): Saga<void> {
   const undoStack: Array<UndoState> = [];
   const redoStack: Array<UndoState> = [];
@@ -133,12 +147,13 @@ export function* collectUndoStates(): Saga<void> {
         }
         previousAction = skeletonUserAction;
       } else if (addBucketToUndoAction) {
-        const { zoomedBucketAddress, bucketData } = addBucketToUndoAction;
+        const { zoomedBucketAddress, bucketData, maybeBucketLoadedPromise } = addBucketToUndoAction;
         pendingCompressions.push(
           yield* fork(
             compressBucketAndAddToUndoBatch,
             zoomedBucketAddress,
             bucketData,
+            maybeBucketLoadedPromise,
             currentVolumeAnnotationBatch,
           ),
         );
@@ -207,6 +222,7 @@ function* getSkeletonTracingToUndoState(
 function* compressBucketAndAddToUndoBatch(
   zoomedBucketAddress: Vector4,
   bucketData: BucketDataArray,
+  maybeBucketLoadedPromise: MaybeBucketLoadedPromise,
   undoBatch: VolumeAnnotationBatch,
 ): Saga<void> {
   const bucketDataAsByteArray = new Uint8Array(
@@ -216,7 +232,23 @@ function* compressBucketAndAddToUndoBatch(
   );
   const compressedBucketData = yield* call(byteArrayToLz4Array, bucketDataAsByteArray, true);
   if (compressedBucketData != null) {
-    undoBatch.push({ zoomedBucketAddress, data: compressedBucketData });
+    const volumeUndoPart: UndoBucket = {
+      zoomedBucketAddress,
+      data: compressedBucketData,
+      maybeBucketLoadedPromise,
+    };
+    if (maybeBucketLoadedPromise != null) {
+      maybeBucketLoadedPromise.then(async backendBucketData => {
+        const backendDataAsByteArray = new Uint8Array(
+          backendBucketData.buffer,
+          backendBucketData.byteOffset,
+          backendBucketData.byteLength,
+        );
+        const compressedBackendData = await byteArrayToLz4Array(backendDataAsByteArray, true);
+        volumeUndoPart.backendData = compressedBackendData;
+      });
+    }
+    undoBatch.push(volumeUndoPart);
   }
 }
 
@@ -274,6 +306,15 @@ function* applyStateOfStack(
   }
 }
 
+function mergeDataWithBackendDataInPlace(
+  originalData: BucketDataArray,
+  backendData: BucketDataArray,
+) {
+  for (let i = 0; i < originalData.length; ++i) {
+    originalData[i] = originalData[i] || backendData[i];
+  }
+}
+
 function* applyAndGetRevertingVolumeBatch(
   volumeAnnotationBatch: VolumeAnnotationBatch,
 ): Saga<VolumeUndoState> {
@@ -283,19 +324,51 @@ function* applyAndGetRevertingVolumeBatch(
   }
   const { cube } = segmentationLayer;
   const allCompressedBucketsOfCurrentState: VolumeAnnotationBatch = [];
-  for (const { zoomedBucketAddress, data: compressedBucketData } of volumeAnnotationBatch) {
+  for (const volumeUndoBucket of volumeAnnotationBatch) {
+    const {
+      zoomedBucketAddress,
+      data: compressedBucketData,
+      backendData: compressedBackendData,
+    } = volumeUndoBucket;
+    let { maybeBucketLoadedPromise } = volumeUndoBucket;
     const bucket = cube.getOrCreateBucket(zoomedBucketAddress);
     if (bucket.type === "null") {
       continue;
     }
     const bucketData = bucket.getData();
+    if (compressedBackendData != null) {
+      // If the backend data for the bucket has been fetched in the meantime,
+      // we can first merge the data with the current data and then add this to the undo batch.
+      const decompressedBackendData = yield* call(
+        byteArrayToLz4Array,
+        compressedBackendData,
+        false,
+      );
+      if (decompressedBackendData) {
+        mergeDataWithBackendDataInPlace(bucketData, decompressedBackendData);
+      }
+      maybeBucketLoadedPromise = null;
+    }
     yield* call(
       compressBucketAndAddToUndoBatch,
       zoomedBucketAddress,
       bucketData,
+      maybeBucketLoadedPromise,
       allCompressedBucketsOfCurrentState,
     );
-    const decompressedBucketData = yield* call(byteArrayToLz4Array, compressedBucketData, false);
+    let decompressedBucketData = null;
+    if (compressedBackendData != null) {
+      let decompressedBackendData;
+      [decompressedBucketData, decompressedBackendData] = yield _all([
+        _call(byteArrayToLz4Array, compressedBucketData, false),
+        _call(byteArrayToLz4Array, compressedBackendData, false),
+      ]);
+      if (decompressedBucketData && decompressedBackendData) {
+        mergeDataWithBackendDataInPlace(decompressedBucketData, decompressedBackendData);
+      }
+    } else {
+      decompressedBucketData = yield* call(byteArrayToLz4Array, compressedBucketData, false);
+    }
     if (decompressedBucketData) {
       // Set the new bucket data to add the bucket directly to the pushqueue.
       cube.setBucketData(zoomedBucketAddress, decompressedBucketData);

frontend/javascripts/oxalis/model/volumetracing/volume_annotation_sampling.js

@@ -332,7 +332,7 @@ export function applyVoxelMap(
       if (bucket.type === "null") {
         continue;
       }
-      const data = bucket.getOrCreateData();
+      const { data } = bucket.getOrCreateData();
       for (let firstDim = 0; firstDim < constants.BUCKET_WIDTH; firstDim++) {
         for (let secondDim = 0; secondDim < constants.BUCKET_WIDTH; secondDim++) {
           if (voxelMap[firstDim * constants.BUCKET_WIDTH + secondDim] === 1) {

frontend/javascripts/test/model/volumetracing/volume_annotation_sampling.spec.js

@@ -472,7 +472,7 @@ test("A labeledVoxelMap should be applied correctly", t => {
     expectedBucketData[addr] = 1;
   });
   applyVoxelMap(labeledVoxelsMap, cube, 1, get3DAddress, 1, 2, true);
-  const labeledBucketData = bucket.getOrCreateData();
+  const { data: labeledBucketData } = bucket.getOrCreateData();
   for (let firstDim = 0; firstDim < Constants.BUCKET_WIDTH; firstDim++) {
     for (let secondDim = 0; secondDim < Constants.BUCKET_WIDTH; secondDim++) {
       const addr = cube.getVoxelIndex([firstDim, secondDim, 5], 0);