Implement undo functionality enhancements for SAM2 points tool

src/wwwroot/js/genpage/helpers/image_editor.js

@@ -320,6 +320,9 @@ class ImageEditorHistoryEntry {
             // TODO: Reinsert at proper index
             this.editor.addLayer(this.data.layer, true);
         }
+        if (this.data.onUndo) {
+            this.data.onUndo();
+        }
     }
 }
 
@@ -466,6 +469,7 @@ class ImageEditor {
 
     undoOnce() {
         if (this.editHistory.length > 0) {
+            this.activeTool.onBeforeHistoryUndo();
             let entry = this.editHistory.pop();
             entry.undo();
             this.redraw();
@@ -987,6 +991,7 @@ class ImageEditor {
         this.realHeight = img.naturalHeight;
         if (this.tools['sam2points']) {
             this.tools['sam2points'].layerPoints = new Map();
+            this.tools['sam2points'].lastAppliedPoints = { positive: [], negative: [] };
         }
         if (this.tools['sam2bbox']) {
             this.tools['sam2bbox'].bboxStartX = null;

src/wwwroot/js/genpage/helpers/image_editor_tools.js

@@ -101,6 +101,9 @@ class ImageEditorTool {
         return false;
     }
 
+    onBeforeHistoryUndo() {
+    }
+
     onLayerChanged(oldLayer, newLayer) {
         if (this.isMaskOnly) {
             let isMask = newLayer && newLayer.isMask;
@@ -1525,6 +1528,34 @@ class ImageEditorToolSam2Points extends ImageEditorToolSam2Base {
         // TODO: This map is a pretty iffy way to do things, probably stray persistence.
         this.layerPoints = new Map();
         this.pendingMaskUpdate = false;
+        this.lastAppliedPoints = { positive: [], negative: [] };
+    }
+
+    flushPointsUndoHistory() {
+        for (let entry of this.editor.editHistory) {
+            delete entry.data.onUndo;
+        }
+    }
+
+    setInactive() {
+        super.setInactive();
+        this.layerPoints = new Map();
+        this.lastAppliedPoints = { positive: [], negative: [] };
+        this.activeRequestId = ++this.requestSerial;
+        this.maskRequestInFlight = false;
+        this.pendingMaskUpdate = false;
+        this.flushPointsUndoHistory();
+    }
+
+    onBeforeHistoryUndo() {
+        this.activeRequestId = ++this.requestSerial;
+        this.maskRequestInFlight = false;
+        this.pendingMaskUpdate = false;
+    }
+
+    onLayerChanged(oldLayer, newLayer) {
+        super.onLayerChanged(oldLayer, newLayer);
+        this.lastAppliedPoints = { positive: [], negative: [] };
     }
 
     getActivePoints() {
@@ -1557,7 +1588,9 @@ class ImageEditorToolSam2Points extends ImageEditorToolSam2Base {
         let points = this.getActivePoints();
         points.positive = [];
         points.negative = [];
+        this.lastAppliedPoints = { positive: [], negative: [] };
         this.clearMaskAndEndRequest();
+        this.flushPointsUndoHistory();
     }
 
     drawPoint(ctx, x, y, fillColor, showX) {
@@ -1689,6 +1722,8 @@ class ImageEditorToolSam2Points extends ImageEditorToolSam2Base {
         if (points.negative.length > 0) {
             genData['samnegativepoints'] = JSON.stringify(points.negative.map(p => ({ x: p.x - offX, y: p.y - offY })));
         }
+        let previousPoints = { positive: [...this.lastAppliedPoints.positive], negative: [...this.lastAppliedPoints.negative] };
+        let thisRequestPoints = { positive: [...points.positive], negative: [...points.negative] };
         makeWSRequestT2I('GenerateText2ImageWS', genData, data => {
             if (requestId != this.activeRequestId || !data.image) {
                 return;
@@ -1703,6 +1738,16 @@ class ImageEditorToolSam2Points extends ImageEditorToolSam2Base {
                     return;
                 }
                 this.applyMaskResult(newImg);
+                let maskLayer = this.editor.activeLayer;
+                let history = this.editor.editHistory;
+                if (history.length > 0) {
+                    let capturedPrevious = previousPoints;
+                    history.at(-1).data.onUndo = () => {
+                        this.layerPoints.set(maskLayer.id, { positive: [...capturedPrevious.positive], negative: [...capturedPrevious.negative] });
+                        this.lastAppliedPoints = capturedPrevious;
+                    };
+                }
+                this.lastAppliedPoints = thisRequestPoints;
                 this.editor.redraw();
                 this.finishMaskUpdate(requestId);
             };