remove background: frontend UX

crates/desktop/src-tauri/src/endpoints/image_endpoint.rs

@@ -32,6 +32,7 @@ pub async fn infer_image(
   app_data_root: State<'_, AppDataRoot>,
   app: AppHandle,
 ) -> Result<String, String> {
+  info!("infer_image endpoint called.");
 
   let prompt = get_prompt_or_fallback(prompt);
 

crates/desktop/src-tauri/src/endpoints/remove_background_endpoint.rs

@@ -6,7 +6,7 @@ use crate::state::app_dir::AppDataRoot;
 use crate::utils::image::decode_base64_image::decode_base64_image;
 use crate::utils::image::encode_dynamic_image_base64_png::encode_dynamic_image_base64_png;
 use image::{imageops, DynamicImage};
-use log::error;
+use log::{error, info};
 use ndarray::{Array, CowArray};
 use ort::Value;
 use tauri::{AppHandle, State};
@@ -19,6 +19,7 @@ pub async fn remove_background(
   app_data_root: State<'_, AppDataRoot>,
   app: AppHandle,
 ) -> Result<String, String> {
+  info!("remove_background endpoint called.");
 
   let image = decode_base64_image(image)
     .map_err(|err| format!("Couldn't hydrate image from base64: {}", err))?;

frontend/realtime/src/KonvaApp/Engine.ts

@@ -1,35 +1,36 @@
 import Konva from "konva";
 
-import { uiAccess, uiEvents } from "~/signals";
-import { ShapeNode } from "./Nodes";
-import { PaintNode } from "./Nodes/PaintNode";
-import { PreviewCopyNode } from "./Nodes/PreviewCopy";
-import { UndoStackManager } from "./UndoRedo";
-import { CommandManager, MatteBox, SceneManager } from "./EngineUtitlities";
+import {uiAccess, uiEvents} from "~/signals";
+import {ImageNode, ShapeNode, ShapeType, TextNode, VideoNode} from "./Nodes";
+import {PaintNode} from "./Nodes/PaintNode";
+import {PreviewCopyNode} from "./Nodes/PreviewCopy";
+import {UndoStackManager} from "./UndoRedo";
+import {CommandManager, MatteBox, SceneManager} from "./EngineUtitlities";
 import {
-  NodesManager,
   NodeIsolator,
-  NodeTransformer,
+  NodesManager,
   NodesTranslationEventDetails,
   NodeTransformationEventDetails,
+  NodeTransformer,
   SelectionManager,
   SelectionManagerEvents,
   SelectorSquare,
 } from "./NodesManagers";
-import { ImageNode, VideoNode, TextNode } from "./Nodes";
-import { EngineOptions, TextNodeData, VideoNodeData } from "./types";
+import {EngineOptions, TextNodeData, VideoNodeData} from "./types";
 
-import { AppModes, VideoResolutions } from "./constants";
-import { ToolbarMainButtonNames } from "~/components/features/ToolbarMain/enum";
+import {AppModes, VideoResolutions} from "./constants";
+import {ToolbarMainButtonNames} from "~/components/features/ToolbarMain/enum";
 
-import { ToolbarNodeButtonNames } from "~/components/features/ToolbarNode/enums";
-import { NavigateFunction } from "react-router-dom";
-import { LoadingVideosProvider } from "./EngineUtitlities/LoadingVideosProvider";
+import {ToolbarNodeButtonNames} from "~/components/features/ToolbarNode/enums";
+import {NavigateFunction} from "react-router-dom";
+import {LoadingVideosProvider} from "./EngineUtitlities/LoadingVideosProvider";
 
-import { VideoExtractionHandler } from "./EngineUtitlities/VideoExtractionHandler/VideoExtractionHandler";
-import { RealTimeDrawEngine } from "./RenderingPrimitives/RealTimeDrawEngine";
-import { NodeColor } from "~/signals/uiEvents/toolbarNode";
-import { ShapeType } from "./Nodes";
+import {VideoExtractionHandler} from "./EngineUtitlities/VideoExtractionHandler/VideoExtractionHandler";
+import {RealTimeDrawEngine} from "./RenderingPrimitives/RealTimeDrawEngine";
+import {NodeColor} from "~/signals/uiEvents/toolbarNode";
+import {invoke} from "@tauri-apps/api/core";
+import {DecodeBase64ToImage} from "~/utilities/DecodeBase64ToImage.ts";
+import {EncodeImageBitmapToBase64} from "~/utilities/EncodeImageBitmapToBase64.ts";
 
 
 export interface RenderingOptions {
@@ -326,6 +327,9 @@ export class Engine {
     uiEvents.toolbarNode.MOVE_LAYER_UP.onClick(() =>
       this.commandManager.moveNodesUp(),
     );
+    uiEvents.toolbarNode.REMOVE_BACKGROUND.onClick(() =>
+      this.removeBackground(),
+    );
 
     // Listen to Toolbar Main
     uiEvents.toolbarMain.UNDO.onClick(() => {
@@ -445,6 +449,38 @@ export class Engine {
     });
   }
 
+  removeBackground() {
+    const selectedNodeSet = this.selectionManager.getSelectedNodes();
+
+    if (selectedNodeSet.size !== 1) {
+      return; // Can only remove background from a single image at a time.
+    }
+
+    const [ selectedNode ] = selectedNodeSet;
+
+    const removeBgInvocation
+        = async (selectedNode: ImageNode, callback: (existingImageNode: ImageNode, createdImage: ImageBitmap) => void) => {
+
+      const image : HTMLImageElement = selectedNode.kNode.image();
+      const base64Bitmap = await EncodeImageBitmapToBase64(image);
+      const base64BitmapResponse = await invoke("remove_background", {
+        image: base64Bitmap,
+      });
+
+      const createdImage = await DecodeBase64ToImage(base64BitmapResponse as string);
+
+      callback(selectedNode, createdImage);
+    };
+
+    const _promise = removeBgInvocation(selectedNode, (selectedNode: ImageNode, createdImage: ImageBitmap) => {
+      if (!this.nodesManager.hasNode(selectedNode)) {
+        return;
+      }
+      this.commandManager.deleteSpecificNodes([selectedNode]);
+      this.addImageFromImageBitmap(createdImage);
+    })
+  }
+
   disableAllButtons() {
     const buttonNames = Object.values(ToolbarNodeButtonNames);
     for (const name of buttonNames) {
@@ -624,6 +660,23 @@ export class Engine {
 
     this.commandManager.createNode(imageNode);
   }
+
+  public addImageFromImageBitmap(imageBitmap: ImageBitmap) {
+    const imageNode = new ImageNode({
+      mediaLayerRef: this.mediaLayer,
+      canvasPosition: this.realTimeDrawEngine.captureCanvas.position(),
+      canvasSize: this.realTimeDrawEngine.captureCanvas.size(),
+      imageBitmap: imageBitmap,
+      selectionManagerRef: this.selectionManager,
+      loaded: async () => {
+        await this.realTimeDrawEngine.render();
+        this.setAppMode(AppModes.SELECT);
+      },
+    });
+
+    this.commandManager.createNode(imageNode);
+  }
+
   public addPaintNode(
     canvas: HTMLCanvasElement,
     lineBounds: {

frontend/realtime/src/KonvaApp/EngineUtitlities/CommandManager.ts

@@ -78,6 +78,20 @@ export class CommandManager {
     this.undoStackManagerRef.executeCommand(command);
     this.renderEngineRef.render();
   }
+
+  deleteSpecificNodes(nodes: MediaNode[]) {
+    const command = new DeleteCommand({
+      nodes: nodes,
+      mediaLayerRef: this.mediaLayerRef,
+      nodesManagerRef: this.nodesManagerRef,
+      nodeTransformerRef: this.nodeTransformerRef,
+      selectionManagerRef: this.selectionManagerRef,
+      renderEngineRef: this.renderEngineRef,
+    });
+    this.undoStackManagerRef.executeCommand(command);
+    this.renderEngineRef.render();
+  }
+
   toggleLockNodes() {
     const nodes = this.selectionManagerRef.getSelectedNodes();
     const node = nodes.values().next().value;

frontend/realtime/src/KonvaApp/Nodes/ImageNode.ts

@@ -12,6 +12,7 @@ interface ImageNodeContructor {
   canvasPosition: Position;
   canvasSize: Size;
   imageFile?: File;
+  imageBitmap?: ImageBitmap;
   mediaFileToken?: string;
   mediaFileUrl?: string;
   transform?: TransformationData;
@@ -34,6 +35,7 @@ export class ImageNode extends NetworkedNode {
     mediaLayerRef,
     selectionManagerRef,
     loaded,
+    imageBitmap,
   }: ImageNodeContructor) {
     // kNodes need to be created first to guaruntee
     // that it is not undefined in parent's context
@@ -73,6 +75,14 @@ export class ImageNode extends NetworkedNode {
       this.loadImageFromUrl(mediaFileUrl);
       return;
     }
+    if (imageBitmap) {
+      this.loadImageFromImageBitmap({
+        imageBitmap: imageBitmap,
+        maxSize: canvasSize,
+        refPosition: canvasPosition,
+      });
+      return;
+    }
     console.log("image node creation is fucked");
   }
   private loadImageFromFile({
@@ -136,6 +146,40 @@ export class ImageNode extends NetworkedNode {
     this.listenToBaseKNode();
   }
 
+  private loadImageFromImageBitmap({
+    imageBitmap,
+    maxSize,
+    refPosition,
+  }: {
+    imageBitmap: ImageBitmap;
+    maxSize: Size;
+    refPosition: Position;
+  }) {
+    this.imageSize = {
+      width: imageBitmap.width,
+      height: imageBitmap.height,
+    };
+    const adjustedSize = NodeUtilities.adjustNodeSizeToCanvas({
+      componentSize: this.imageSize,
+      maxSize: maxSize, // TODO
+    });
+    const centerPosition = NodeUtilities.positionNodeOnCanvasCenter({
+      canvasOffset: refPosition, // TODO
+      componentSize: adjustedSize, // TODO
+      maxSize: maxSize, // TODO
+    });
+    this.kNode.image(imageBitmap);
+    this.kNode.setSize(adjustedSize);
+    this.kNode.setPosition(centerPosition);
+    this.kNode.moveToTop();
+    this.kNode.fill(transparent);
+    this.listenToBaseKNode();
+    this.mediaLayerRef.draw();
+    //this.uploadImage(imageFile);
+    this.loaded();
+    this.listenToBaseKNode();
+  }
+
   private async uploadImage(imageFile: File) {
     //this.setProgress({ progress: 10, status: UploadStatus.UPLOADING });
 

frontend/realtime/src/KonvaApp/NodesManagers/NodesManager.ts

@@ -44,6 +44,10 @@ export class NodesManager {
     // this.printNodesAndZIndices();
   }
 
+  public hasNode(node: MediaNode): boolean {
+    return this.allNodes.has(node);
+  }
+
   /*************************
    *  For Debugging use
    */

frontend/realtime/src/KonvaApp/NodesManagers/ToolbarNodeUtilities.ts

@@ -94,7 +94,8 @@ export function getTextNodeButtonStates(
         buttonName === ButtonNames.AI_STYLIZE ||
         buttonName === ButtonNames.SEGMENTATION ||
         buttonName === ButtonNames.DOWNLOAD ||
-        buttonName === ButtonNames.CHROMA
+        buttonName === ButtonNames.CHROMA ||
+        buttonName === ButtonNames.REMOVE_BACKGROUND
       ) {
         buttonStates[buttonName] = {
           disabled: true,
@@ -132,7 +133,8 @@ export function getVideoNodeButtonStates(
       }
       // soon to come feature is disabled
       if (buttonName === ButtonNames.AI_STYLIZE ||
-        buttonName === ButtonNames.COLOR
+        buttonName === ButtonNames.COLOR ||
+        buttonName === ButtonNames.REMOVE_BACKGROUND
       ) {
         buttonStates[buttonName] = {
           disabled: true,
@@ -174,7 +176,8 @@ export function getShapeNodeButtonStates(
         buttonName === ButtonNames.AI_STYLIZE ||
         buttonName === ButtonNames.SEGMENTATION ||
         buttonName === ButtonNames.DOWNLOAD ||
-        buttonName === ButtonNames.CHROMA
+        buttonName === ButtonNames.CHROMA ||
+        buttonName === ButtonNames.REMOVE_BACKGROUND
       ) {
         buttonStates[buttonName] = {
           disabled: true,
@@ -218,7 +221,8 @@ export function getMultiSelectButtonStates(
         buttonName === ButtonNames.SEGMENTATION ||
         buttonName === ButtonNames.DOWNLOAD ||
         buttonName === ButtonNames.CHROMA ||
-        buttonName === ButtonNames.COLOR
+        buttonName === ButtonNames.COLOR ||
+        buttonName === ButtonNames.REMOVE_BACKGROUND
       ) {
         buttonStates[buttonName] = {
           disabled: true,

frontend/realtime/src/KonvaApp/RenderingPrimitives/RealTimeDrawEngine.ts

@@ -25,6 +25,7 @@ import {
   isLoadingVisible,
   loadingProgress,
 } from "~/signals/uiEvents/loadingIndicator";
+import {DecodeBase64ToImage} from "~/utilities/DecodeBase64ToImage.ts";
 
 interface ServerSettings {
   model_path: string;
@@ -285,7 +286,7 @@ export class RealTimeDrawEngine {
       this.client.onResult(async (response) => {
         try {
           // Convert the base64 image to ImageBitmap and update the preview
-          const bitmap = await this.base64ToImageBitmap(response.image);
+          const bitmap = await DecodeBase64ToImage(response.image);
           this.outputBitmap = bitmap;
           this.previewCanvas.image(bitmap);
           this.mediaLayerRef.batchDraw();
@@ -895,35 +896,6 @@ export class RealTimeDrawEngine {
     return base64String.split(",")[1];
   }
 
-  private async base64ToImageBitmap(
-    base64String: string,
-  ): Promise<ImageBitmap> {
-    // Create an image element
-    const img = document.createElement("img");
-
-    // Convert base64 to data URL if it doesn't include the prefix
-    const dataUrl = base64String.startsWith("data:")
-      ? base64String
-      : `data:image/png;base64,${base64String}`;
-
-    // Create a promise to handle the image loading
-    return new Promise((resolve, reject) => {
-      img.onload = async () => {
-        try {
-          const bitmap = await createImageBitmap(img);
-          resolve(bitmap);
-        } catch (error) {
-          reject(error);
-        }
-      };
-
-      img.onerror = () => reject(new Error("Failed to load image"));
-
-      // Set the source to trigger loading
-      img.src = dataUrl;
-    });
-  }
-
   public async saveOutput() {
     if (!this.outputBitmap) {
       console.error("No output bitmap available to save");
@@ -1007,7 +979,7 @@ export class RealTimeDrawEngine {
       });
 
       //console.log(base64BitmapResponse);
-      const decoded = await this.base64ToImageBitmap(
+      const decoded = await DecodeBase64ToImage(
         base64BitmapResponse as string,
       );
 

frontend/realtime/src/components/features/ToolbarNode/data.ts

@@ -22,6 +22,11 @@ export const ToolbarNodeButtonData = [
     icon: faPalette,
     tooltip: "Color",
   },
+  {
+    name: ButtonNames.REMOVE_BACKGROUND,
+    icon: faScalpelLineDashed,
+    tooltip: "Remove Background",
+  },
   {
     name: ButtonNames.CHROMA,
     icon: faTransporter,

frontend/realtime/src/components/features/ToolbarNode/enums.ts

@@ -8,4 +8,5 @@ export enum ToolbarNodeButtonNames {
   MOVE_LAYER_DOWN = "MOVE_LAYER_DOWN",
   DOWNLOAD = "DOWNLOAD",
   COLOR = "COLOR",
+  REMOVE_BACKGROUND = "REMOVE_BACKGROUND",
 }

frontend/realtime/src/utilities/DecodeBase64ToImage.ts

@@ -0,0 +1,30 @@
+/**
+ * This decodes Base64 encoded PNG images into an image element.
+ * @param base64String a standard (not "web safe") base64-encoded string
+ */
+export const DecodeBase64ToImage = async (base64String: string) : Promise<ImageBitmap> => {
+  // Create an image element
+  const img = document.createElement("img");
+
+  // Convert base64 to data URL if it doesn't include the prefix
+  const dataUrl = base64String.startsWith("data:")
+    ? base64String
+    : `data:image/png;base64,${base64String}`;
+
+  // Create a promise to handle the image loading
+  return new Promise((resolve, reject) => {
+    img.onload = async () => {
+      try {
+        const bitmap = await createImageBitmap(img);
+        resolve(bitmap);
+      } catch (error) {
+        reject(error);
+      }
+    };
+
+    img.onerror = () => reject(new Error("Failed to load image"));
+
+    // Set the source to trigger loading
+    img.src = dataUrl;
+  });
+}

frontend/realtime/src/utilities/EncodeImageBitmapToBase64.ts

@@ -0,0 +1,22 @@
+
+export const EncodeImageBitmapToBase64 = async (imageBitmap: ImageBitmap): Promise<string> => {
+  // Create a temporary canvas
+  const canvas = document.createElement("canvas");
+  canvas.width = imageBitmap.width;
+  canvas.height = imageBitmap.height;
+
+  // Draw the ImageBitmap onto the canvas
+  const ctx = canvas.getContext("2d");
+  if (!ctx) throw new Error("Failed to get 2D context");
+
+  ctx.drawImage(imageBitmap, 0, 0);
+
+  // Convert to base64
+  const base64String = canvas.toDataURL("image/png");
+
+  // Clean up
+  canvas.remove();
+
+  // Remove the data:image/png;base64, prefix if you want just the base64 string
+  return base64String.split(",")[1];
+}