From f37011ffc637f65cf48cc0ead86a10c6662bc234 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Thu, 20 Mar 2025 01:41:43 -0400 Subject: [PATCH] remove background: frontend UX (#458) --- .../src-tauri/src/endpoints/image_endpoint.rs | 1 + .../endpoints/remove_background_endpoint.rs | 3 +- frontend/realtime/src/KonvaApp/Engine.ts | 91 +++++++++++++++---- .../EngineUtitlities/CommandManager.ts | 14 +++ .../realtime/src/KonvaApp/Nodes/ImageNode.ts | 44 +++++++++ .../KonvaApp/NodesManagers/NodesManager.ts | 4 + .../NodesManagers/ToolbarNodeUtilities.ts | 12 ++- .../RenderingPrimitives/RealTimeDrawEngine.ts | 34 +------ .../components/features/ToolbarNode/data.ts | 5 + .../components/features/ToolbarNode/enums.ts | 1 + .../src/utilities/DecodeBase64ToImage.ts | 30 ++++++ .../utilities/EncodeImageBitmapToBase64.ts | 22 +++++ 12 files changed, 206 insertions(+), 55 deletions(-) create mode 100644 frontend/realtime/src/utilities/DecodeBase64ToImage.ts create mode 100644 frontend/realtime/src/utilities/EncodeImageBitmapToBase64.ts diff --git a/crates/desktop/src-tauri/src/endpoints/image_endpoint.rs b/crates/desktop/src-tauri/src/endpoints/image_endpoint.rs index fa0db4e540..9709736ba0 100644 --- a/crates/desktop/src-tauri/src/endpoints/image_endpoint.rs +++ b/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 { + info!("infer_image endpoint called."); let prompt = get_prompt_or_fallback(prompt); diff --git a/crates/desktop/src-tauri/src/endpoints/remove_background_endpoint.rs b/crates/desktop/src-tauri/src/endpoints/remove_background_endpoint.rs index 8d8b7aa5c9..487315b866 100644 --- a/crates/desktop/src-tauri/src/endpoints/remove_background_endpoint.rs +++ b/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 { + info!("remove_background endpoint called."); let image = decode_base64_image(image) .map_err(|err| format!("Couldn't hydrate image from base64: {}", err))?; diff --git a/frontend/realtime/src/KonvaApp/Engine.ts b/frontend/realtime/src/KonvaApp/Engine.ts index ef2814b226..91b3cbdb42 100644 --- a/frontend/realtime/src/KonvaApp/Engine.ts +++ b/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: { diff --git a/frontend/realtime/src/KonvaApp/EngineUtitlities/CommandManager.ts b/frontend/realtime/src/KonvaApp/EngineUtitlities/CommandManager.ts index 40aa02b087..5463d4e9e4 100644 --- a/frontend/realtime/src/KonvaApp/EngineUtitlities/CommandManager.ts +++ b/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; diff --git a/frontend/realtime/src/KonvaApp/Nodes/ImageNode.ts b/frontend/realtime/src/KonvaApp/Nodes/ImageNode.ts index fc57c971c2..a81cdf7301 100644 --- a/frontend/realtime/src/KonvaApp/Nodes/ImageNode.ts +++ b/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 }); diff --git a/frontend/realtime/src/KonvaApp/NodesManagers/NodesManager.ts b/frontend/realtime/src/KonvaApp/NodesManagers/NodesManager.ts index 5fec9128fd..f86afc2018 100644 --- a/frontend/realtime/src/KonvaApp/NodesManagers/NodesManager.ts +++ b/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 */ diff --git a/frontend/realtime/src/KonvaApp/NodesManagers/ToolbarNodeUtilities.ts b/frontend/realtime/src/KonvaApp/NodesManagers/ToolbarNodeUtilities.ts index 8e09c585be..0c14479ac3 100644 --- a/frontend/realtime/src/KonvaApp/NodesManagers/ToolbarNodeUtilities.ts +++ b/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, diff --git a/frontend/realtime/src/KonvaApp/RenderingPrimitives/RealTimeDrawEngine.ts b/frontend/realtime/src/KonvaApp/RenderingPrimitives/RealTimeDrawEngine.ts index a1f151596c..4d681ec479 100644 --- a/frontend/realtime/src/KonvaApp/RenderingPrimitives/RealTimeDrawEngine.ts +++ b/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 { - // 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, ); diff --git a/frontend/realtime/src/components/features/ToolbarNode/data.ts b/frontend/realtime/src/components/features/ToolbarNode/data.ts index c215e5ed8b..f8bc6a8d71 100644 --- a/frontend/realtime/src/components/features/ToolbarNode/data.ts +++ b/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, diff --git a/frontend/realtime/src/components/features/ToolbarNode/enums.ts b/frontend/realtime/src/components/features/ToolbarNode/enums.ts index e4c559695a..7401d94e92 100644 --- a/frontend/realtime/src/components/features/ToolbarNode/enums.ts +++ b/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", } diff --git a/frontend/realtime/src/utilities/DecodeBase64ToImage.ts b/frontend/realtime/src/utilities/DecodeBase64ToImage.ts new file mode 100644 index 0000000000..b568808afd --- /dev/null +++ b/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 => { + // 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; + }); +} diff --git a/frontend/realtime/src/utilities/EncodeImageBitmapToBase64.ts b/frontend/realtime/src/utilities/EncodeImageBitmapToBase64.ts new file mode 100644 index 0000000000..31f8ed79c3 --- /dev/null +++ b/frontend/realtime/src/utilities/EncodeImageBitmapToBase64.ts @@ -0,0 +1,22 @@ + +export const EncodeImageBitmapToBase64 = async (imageBitmap: ImageBitmap): Promise => { + // 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]; +}