cleaned up image background removal code a bit and added undo

3 files changed, 56 insertions, 48 deletions

diff --git a/src/client/util/Import & Export/ImageUtils.ts b/src/client/util/Import & Export/ImageUtils.ts
index e2eab9010..23102e051 100644
--- a/src/client/util/Import & Export/ImageUtils.ts
+++ b/src/client/util/Import & Export/ImageUtils.ts
@@ -41,42 +41,47 @@ export namespace ImageUtils {
             reader.onerror = e => reject(e);
             reader.readAsDataURL(blob as Blob);
         });
-    };
-    export function createImageDocFromBlob(blob: Blob, options: DocumentOptions, docId: string, nativeWidth: number, nativeHeight: number): Promise<Doc> {
+    }
+    export function createImageDocFromBlob(blob: Blob | undefined, options: DocumentOptions & { _nativeWidth: number; _nativeHeight: number }, filename: string): Promise<Doc> {
         return new Promise((resolve, reject) => {
-            convertImgBlobToDataURL(blob).then(durl => {
-                ClientUtils.convertDataUri(durl as string, docId + '_noBgd').then(url => {
-                    const imageSnapshot = Docs.Create.ImageDocument(url, options);
-                    Doc.SetNativeWidth(imageSnapshot[DocData], nativeWidth);
-                    Doc.SetNativeHeight(imageSnapshot[DocData], nativeHeight);
-                    resolve(imageSnapshot);
-                }).catch(reject);
-            }).catch(reject);
+            if (!blob) return reject('No image blob provided');
+            convertImgBlobToDataURL(blob)
+                .then(durl => {
+                    ClientUtils.convertDataUri(durl as string, filename)
+                        .then(url => {
+                            const imageSnapshot = Docs.Create.ImageDocument(url, options);
+                            Doc.SetNativeWidth(imageSnapshot[DocData], options._nativeWidth);
+                            Doc.SetNativeHeight(imageSnapshot[DocData], options._nativeHeight);
+                            resolve(imageSnapshot);
+                        })
+                        .catch(reject);
+                })
+                .catch(reject);
         });
     }
-    let backgroundRemovalWorker: Worker|undefined;
-    const workerCallbackMap = new Map<string, (result: Doc) => void>();
-    function getBackgroundRemovalWorker()  {
+    let backgroundRemovalWorker: Worker | undefined;
+    const workerCallbackMap = new Map<string, { res: (imgBlob: Blob) => unknown; rej: () => unknown }>();
+    function getBackgroundRemovalWorker() {
         if (backgroundRemovalWorker) return backgroundRemovalWorker;
         backgroundRemovalWorker = new Worker('/image.worker.js', { type: 'module' });
-        backgroundRemovalWorker.onerror = (e: ErrorEvent) => console.error('Worker failed:', e); //  worker.terminate();
+        backgroundRemovalWorker.onerror = (e: ErrorEvent) => {
+            const docId = (e as ErrorEvent & { docId: string }).docId;
+            const map = workerCallbackMap.get(docId);
+            map?.rej();
+            workerCallbackMap.delete(docId); //  worker.terminate();
+        };
         backgroundRemovalWorker.onmessage = async (event: MessageEvent) => {
-            if (event.data.success) {
-                createImageDocFromBlob(event.data.result, event.data.options, event.data.docId, event.data.nativeWidth, event.data.nativeHeight)
-                    .then(imageSnapshot => {
-                        workerCallbackMap.get(event.data.docId)?.(imageSnapshot);
-                        workerCallbackMap.delete(event.data.docId);
-                    })
-            } else {
-                console.error('Error in background removal:', event.data.error);
-            }
-            // worker.terminate();
+            const map = workerCallbackMap.get(event.data.docId);
+            event.data.success ? map?.res(event.data.result) : map?.rej();
+            workerCallbackMap.delete(event.data.docId); // worker.terminate();
         };
         return backgroundRemovalWorker;
-
     }
-    export function removeImgBackground(options:DocumentOptions, cb: (result: Doc) => void, docId: string, nativeWidth:number, nativeHeight:number, docImgPath: string) {
-        workerCallbackMap.set(docId, cb); // Store the callback by docId (or use a unique requestId)
-        getBackgroundRemovalWorker().postMessage({ imagePath: docImgPath, options, docId, nativeWidth, nativeHeight });
-    };
+    export function removeImgBackground(docId: string, imagePath: string) {
+        return new Promise<Blob | undefined>((res, rej) => {
+            if (!imagePath) return rej('No image path provided');
+            workerCallbackMap.set(docId, { res, rej }); // Store the callback by docId (or use a unique requestId)
+            getBackgroundRemovalWorker().postMessage({ imagePath, docId });
+        });
+    }
 }
diff --git a/src/client/views/nodes/ImageBox.tsx b/src/client/views/nodes/ImageBox.tsx
index c4a731078..e62718ec4 100644
--- a/src/client/views/nodes/ImageBox.tsx
+++ b/src/client/views/nodes/ImageBox.tsx
@@ -390,22 +390,25 @@ export class ImageBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
         return cropping;
     };
     removeBackground = () => {
-        const field = ImageCast(this.dataDoc[this.fieldKey]);
-        if (field && this._props.addDocument) {
-            const doc = this.Document;
-            const width = NumCast(doc._width) || 1;
-            const height = NumCast(doc._height);
-            const options ={
-                _nativeWidth: Doc.NativeWidth(doc),
-                _nativeHeight: Doc.NativeHeight(doc),
-                x: NumCast(doc.x) + width,
-                y: NumCast(doc.y),
-                _width: 150,
-                _height: (height / width) * 150,
-                title: 'bgremoved:' + doc.title,
-            }
-            ImageUtils.removeImgBackground(options, this._props.addDocument, this.Document[Id], Doc.NativeWidth(doc), Doc.NativeHeight(doc), this.choosePath(field.url));
-        }
+        const batch = UndoManager.StartBatch('remove image background');
+        ImageUtils.removeImgBackground(this.Document[Id], this.choosePath(ImageCast(this.dataDoc[this.fieldKey])?.url)).then(imgBlob =>
+            ImageUtils.createImageDocFromBlob(
+                imgBlob,
+                {
+                    _nativeWidth: Doc.NativeWidth(this.Document),
+                    _nativeHeight: Doc.NativeHeight(this.Document),
+                    x: NumCast(this.Document.x) + NumCast(this.Document._width),
+                    y: NumCast(this.Document.y),
+                    _width: NumCast(this.Document._width),
+                    _height: (NumCast(this.Document._height) / (NumCast(this.Document._width) || 1)) * NumCast(this.Document._width),
+                    title: 'bgdRemoved:' + this.Document.title,
+                },
+                this.Document[Id] + '_noBgd'
+            ).then(imageSnapshot => {
+                this._props.addDocument?.(imageSnapshot);
+                batch.end();
+            })
+        );
     };
 
     docEditorView = () => {
@@ -648,7 +651,7 @@ export class ImageBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
                         file: (file => {
                             const ext = file ? extname(file) : '';
                             return file?.replace(ext, (this._error ? '_o' : this._curSuffix) + ext);
-                        })(ImageCast(this.Document[this.fieldKey])?.url.href),
+                        })(field.url.href),
                     }).then(text => alert(text));
                 },
                 icon: 'expand-arrows-alt',
@@ -716,7 +719,7 @@ export class ImageBox extends ViewBoxAnnotatableComponent<FieldViewProps>() {
                   }
               );
 
-    choosePath = (url: URL) => {
+    choosePath = (url: URL | undefined) => {
         if (!url?.href) return '';
         const lower = url.href.toLowerCase();
         if (url.protocol === 'data') return url.href;
diff --git a/src/workers/image.worker.ts b/src/workers/image.worker.ts
index c6db44d03..48b4e8585 100644
--- a/src/workers/image.worker.ts
+++ b/src/workers/image.worker.ts
@@ -12,6 +12,6 @@ self.onmessage = async (event: MessageEvent) => {
     } catch (error) {
         // Send the error back to the main thread
         // eslint-disable-next-line @typescript-eslint/no-explicit-any
-        self.postMessage({ success: false, error: (error as any).message });
+        self.postMessage({ success: false, docId, error: (error as any).message });
     }
 };