remove background from images

Cargo.lock

@@ -676,7 +676,9 @@ dependencies = [
  "libloading 0.8.6",
  "log",
  "memory_caching",
+ "ndarray",
  "once_cell",
+ "ort",
  "primitives",
  "rand 0.8.5",
  "reqwest 0.12.12",
@@ -7043,6 +7045,19 @@ dependencies = [
  "tempfile",
 ]
 
+[[package]]
+name = "ndarray"
+version = "0.15.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "adb12d4e967ec485a5f71c6311fe28158e9d6f4bc4a447b474184d0f91a8fa32"
+dependencies = [
+ "matrixmultiply",
+ "num-complex",
+ "num-integer",
+ "num-traits",
+ "rawpointer",
+]
+
 [[package]]
 name = "ndk"
 version = "0.8.0"
@@ -7825,6 +7840,26 @@ dependencies = [
  "hashbrown 0.14.5",
 ]
 
+[[package]]
+name = "ort"
+version = "1.16.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "889dca4c98efa21b1ba54ddb2bde44fd4920d910f492b618351f839d8428d79d"
+dependencies = [
+ "flate2",
+ "half",
+ "lazy_static",
+ "libc",
+ "ndarray",
+ "tar",
+ "thiserror 1.0.63",
+ "tracing",
+ "ureq",
+ "vswhom",
+ "winapi 0.3.9",
+ "zip 0.6.6",
+]
+
 [[package]]
 name = "overload"
 version = "0.1.1"
@@ -12606,9 +12641,9 @@ dependencies = [
 
 [[package]]
 name = "tracing-subscriber"
-version = "0.3.19"
+version = "0.3.18"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "e8189decb5ac0fa7bc8b96b7cb9b2701d60d48805aca84a238004d665fcc4008"
+checksum = "ad0f048c97dbd9faa9b7df56362b8ebcaa52adb06b498c050d2f4e32f90a7a8b"
 dependencies = [
  "nu-ansi-term",
  "sharded-slab",

crates/desktop/src-tauri/Cargo.toml

@@ -58,6 +58,11 @@ tokio = { version = "1.42", features = ["rt", "rt-multi-thread", "fs"] }
 tokio-util = { version = "0.7", features = ["codec"] }
 half = { version = "2.3.1", features = ["num-traits", "use-intrinsics", "rand_distr"] }
 
+# NB: From `rust-background-removal`
+ndarray = "0.15.6"
+ort = { version = "1.16.2", features = ["cuda"] }
+# tracing-subscriber = { version = "0.3", features = [ "env-filter", "fmt" ] }
+
 [features]
 default = ["cuda"]
 cuda = [

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

@@ -3,12 +3,11 @@ use crate::ml::prompt_cache::PromptCache;
 use crate::ml::stable_diffusion::lcm_pipeline::{lcm_pipeline, Args};
 use crate::state::app_config::AppConfig;
 use crate::state::app_dir::AppDataRoot;
-use base64::{engine::general_purpose::STANDARD as BASE64_STANDARD, Engine};
-use bytes::BytesMut;
+use crate::utils::image::decode_base64_image::decode_base64_image;
+use crate::utils::image::encode_rgb_image_base64_png::encode_rgb_image_base64_png;
 use image::imageops::FilterType;
-use image::{DynamicImage, ImageFormat, ImageReader};
+use image::{DynamicImage, RgbImage};
 use log::{error, info};
-use std::io::Cursor;
 use std::path::PathBuf;
 use tauri::{AppHandle, State};
 
@@ -38,9 +37,12 @@ pub async fn infer_image(
 
   info!("Strength: {:?}; Prompt: {}", strength, prompt);
 
-  let image = hydrate_base64_image(image)
+  let image = decode_base64_image(image)
     .map_err(|err| format!("Couldn't hydrate image from base64: {}", err))?;
   
+  // TODO(bt,2025-02-17): Running out of vram with full image buffer size
+  let image = image.resize(512, 512, FilterType::CatmullRom);
+
   let result = do_infer_image(
     &prompt, 
     image, 
@@ -52,11 +54,20 @@ pub async fn infer_image(
     &app_data_root
   ).await;
   
-  if let Err(err) = result.as_deref() {
-    error!("There was an error: {:?}", err);
-  }
-  
-  result
+  let image = match result {
+    Ok(image) => image,
+    Err(err) => {
+      error!("There was an error generating the image: {:?}", err);
+      return Err(format!("There was an error generating the image: {}", err));
+    }
+  };
+
+  let bytes = encode_rgb_image_base64_png(image)
+    .map_err(|err| format!("failure to encode image: {:?}", err))?;
+
+  info!("Inference successful; image converted to base64, serving back to browser...");
+
+  Ok(bytes)
 }
 
 async fn do_infer_image(
@@ -68,7 +79,7 @@ async fn do_infer_image(
   prompt_cache: State<'_, PromptCache>,
   app: AppHandle,
   app_data_root: &AppDataRoot,
-) -> Result<String, String> {
+) -> Result<RgbImage, String> {
 
   let args = Args {
     image: &image,
@@ -84,21 +95,10 @@ async fn do_infer_image(
     use_flash_attn: true,
   };
 
-  match lcm_pipeline(args) {
-    Ok(image) => {
-      let mut bytes = Vec::with_capacity(1024*1024);
-      
-      image.write_to(&mut Cursor::new(&mut bytes), ImageFormat::Png)
-        .map_err(|err| format!("failure to encode image: {:?}", err))?;
+  let image = lcm_pipeline(args)
+    .map_err(|err| format!("failure to encode image: {:?}", err))?;
 
-      let bytes = BASE64_STANDARD.encode(bytes);
-
-      info!("Inference successful; image converted to base64, serving back to browser...");
-
-      Ok(bytes)
-    },
-    Err(e) => Err(format!("Failed to generate image: {}", e)),
-  }
+  Ok(image)
 }
 
 fn get_prompt_or_fallback(user_prompt: Option<String>) -> String {
@@ -122,29 +122,3 @@ fn get_prompt_or_fallback(user_prompt: Option<String>) -> String {
     .trim()
     .to_string()
 }
-
-fn hydrate_base64_image(base64_image: &str) -> anyhow::Result<DynamicImage> {
-  let bytes = BASE64_STANDARD.decode(base64_image)?;
-
-  let image = ImageReader::new(Cursor::new(bytes))
-    .with_guessed_format()?
-    .decode()?;
-
-  // TODO(bt,2025-02-17): Running out of vram with full image buffer size
-  let image = image.resize(512, 512, FilterType::CatmullRom);
-
-  Ok(image)
-}
-
-fn float_to_u8(value: f64) -> u8 {
-  const MAX : f64 = u8::MAX as f64;
-  const MIN : f64 = 0f64;
-  let value = value.round();
-  if value >= MAX {
-    u8::MAX
-  } else if value <= MIN {
-    0
-  } else {
-    value as u8
-  }
-}

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

@@ -1,4 +1,5 @@
 pub mod download_models;
 pub mod flip_image;
 pub mod image_endpoint;
+pub mod remove_background_endpoint;
 pub mod test_counter;

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

@@ -0,0 +1,90 @@
+//! Adapted from https://github.com/dnanhkhoa/rust-background-removal
+
+use crate::ml::background_removal::onnx_session::onnx_session;
+use crate::ml::model_cache::ModelCache;
+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 ndarray::{Array, CowArray};
+use ort::Value;
+use tauri::{AppHandle, State};
+
+/// This handler removes the background from an image.
+#[tauri::command]
+pub async fn remove_background(
+  image: &str,
+  model_cache: State<'_, ModelCache>,
+  app_data_root: State<'_, AppDataRoot>,
+  app: AppHandle,
+) -> Result<String, String> {
+
+  let image = decode_base64_image(image)
+    .map_err(|err| format!("Couldn't hydrate image from base64: {}", err))?;
+
+  // NB: Real time needed cuDNN 9, but this uses cuDNN 8:
+  // sudo apt-get install libcudnn8-dev libcudnn8
+  // update-alternatives: warning: forcing reinstallation of alternative /usr/include/x86_64-linux-gnu/cudnn_v9.h because link group libcudnn is broken
+  // update-alternatives: using /usr/include/x86_64-linux-gnu/cudnn_v8.h to provide /usr/include/cudnn.h (libcudnn) in manual mode
+  let model_path = app_data_root.weights_dir().path().join("medium.onnx"); // TODO
+  let session = onnx_session(model_path.to_str().unwrap()) // TODO
+    .map_err(|err| format!("failure to start onnx session: {:?}", err))?;
+
+  let image = process_dynamic_image(&session, image)
+    .map_err(|err| {
+      error!("Couldn't remove background from image: {}", err);
+      format!("Couldn't remove background from image: {}", err)
+    })?;
+
+  let image = encode_dynamic_image_base64_png(image)
+    .map_err(|err| format!("failure to encode image: {:?}", err))?;
+
+  Ok(image)
+}
+
+fn process_dynamic_image(
+  session: &ort::Session,
+  dynamic_img: DynamicImage,
+) -> Result<DynamicImage, anyhow::Error> {
+  let input_shape = session.inputs[0]
+    .dimensions()
+    .map(|dim| dim.unwrap())
+    .collect::<Vec<usize>>();
+  let input_img = dynamic_img.into_rgba8();
+  let scaling_factor = f32::min(
+    1., // Avoid upscaling
+    f32::min(
+      input_shape[3] as f32 / input_img.width() as f32, // Width ratio
+      input_shape[2] as f32 / input_img.height() as f32, // Height ratio
+    ),
+  );
+  let mut resized_img = imageops::resize(
+    &input_img,
+    input_shape[3] as u32,
+    input_shape[2] as u32,
+    imageops::FilterType::Triangle,
+  );
+  let input_tensor = CowArray::from(
+    Array::from_shape_fn(input_shape, |indices| {
+      let mean = 128.;
+      let std = 256.;
+
+      (resized_img[(indices[3] as u32, indices[2] as u32)][indices[1]] as f32 - mean) / std
+    })
+      .into_dyn(),
+  );
+  let inputs = vec![Value::from_array(session.allocator(), &input_tensor)?];
+  let outputs = session.run(inputs)?;
+  let output_tensor = outputs[0].try_extract::<f32>()?;
+  for (indices, alpha) in output_tensor.view().indexed_iter() {
+    resized_img[(indices[3] as u32, indices[2] as u32)][3] = (alpha * 255.) as u8;
+  }
+  let output_img = imageops::resize(
+    &resized_img,
+    (input_img.width() as f32 * scaling_factor) as u32,
+    (input_img.height() as f32 * scaling_factor) as u32,
+    imageops::FilterType::Triangle,
+  );
+  Ok(DynamicImage::ImageRgba8(output_img))
+}

crates/desktop/src-tauri/src/lib.rs

@@ -4,9 +4,11 @@ pub mod ml;
 pub mod state;
 pub mod threads;
 pub mod transfer;
+pub mod utils;
 
 use crate::endpoints::download_models::download_models;
 use crate::endpoints::image_endpoint::infer_image;
+use crate::endpoints::remove_background_endpoint::remove_background;
 use crate::endpoints::test_counter::test_counter;
 use crate::ml::model_cache::ModelCache;
 use crate::ml::prompt_cache::PromptCache;
@@ -60,9 +62,10 @@ pub fn run() {
     .manage(model_cache)
     .manage(app_data_root)
     .invoke_handler(tauri::generate_handler![
-      test_counter, 
-      infer_image,
       download_models,
+      infer_image,
+      remove_background,
+      test_counter, 
     ])
     .run(tauri::generate_context!())
     .expect("error while running tauri application");

crates/desktop/src-tauri/src/ml/background_removal/mod.rs

@@ -0,0 +1 @@
+pub mod onnx_session;

crates/desktop/src-tauri/src/ml/background_removal/onnx_session.rs

@@ -0,0 +1,18 @@
+//! Adapted from https://github.com/dnanhkhoa/rust-background-removal/blob/main/src/onnx.rs
+
+use ort::{Environment, ExecutionProvider, GraphOptimizationLevel, SessionBuilder};
+
+pub fn onnx_session(onnx_model_file: &str) -> Result<ort::Session, anyhow::Error> {
+  let environment = Environment::default().into_arc();
+  let session = SessionBuilder::new(&environment)?
+    .with_optimization_level(GraphOptimizationLevel::Level1)? // Configure model optimization level
+    .with_intra_threads(1)? // Configure the number of threads used for inference
+    .with_execution_providers([
+      // Configure execution providers (e.g., CUDA, CoreML, CPU)
+      ExecutionProvider::CUDA(Default::default()),
+      ExecutionProvider::CoreML(Default::default()),
+      ExecutionProvider::CPU(Default::default()),
+    ])?
+    .with_model_from_file(onnx_model_file)?;
+  Ok(session)
+}

crates/desktop/src-tauri/src/ml/mod.rs

@@ -1,3 +1,4 @@
+pub mod background_removal;
 pub mod image;
 pub mod model_cache;
 pub mod model_file;

crates/desktop/src-tauri/src/utils/image/decode_base64_image.rs

@@ -0,0 +1,14 @@
+use base64::prelude::BASE64_STANDARD;
+use base64::Engine;
+use image::{DynamicImage, ImageReader};
+use std::io::Cursor;
+
+pub fn decode_base64_image(base64_image: &str) -> anyhow::Result<DynamicImage> {
+  let bytes = BASE64_STANDARD.decode(base64_image)?;
+
+  let image = ImageReader::new(Cursor::new(bytes))
+    .with_guessed_format()?
+    .decode()?;
+
+  Ok(image)
+}

crates/desktop/src-tauri/src/utils/image/encode_dynamic_image_base64_png.rs

@@ -0,0 +1,14 @@
+use base64::prelude::BASE64_STANDARD;
+use base64::Engine;
+use image::{DynamicImage, ImageFormat};
+use std::io::Cursor;
+
+pub fn encode_dynamic_image_base64_png(image: DynamicImage) -> anyhow::Result<String> {
+  let mut bytes = Vec::with_capacity(1024*1024);
+
+  image.write_to(&mut Cursor::new(&mut bytes), ImageFormat::Png)?;
+
+  let bytes = BASE64_STANDARD.encode(bytes);
+
+  Ok(bytes)
+}

crates/desktop/src-tauri/src/utils/image/encode_rgb_image_base64_png.rs

@@ -0,0 +1,17 @@
+use crate::ml::image::tensor_to_image_buffer::RgbImage;
+use base64::prelude::BASE64_STANDARD;
+use base64::Engine;
+use image::ImageFormat;
+use std::io::Cursor;
+
+pub fn encode_rgb_image_base64_png(
+  image: RgbImage
+) -> anyhow::Result<String> {
+  let mut bytes = Vec::with_capacity(1024*1024);
+
+  image.write_to(&mut Cursor::new(&mut bytes), ImageFormat::Png)?;
+
+  let bytes = BASE64_STANDARD.encode(bytes);
+
+  Ok(bytes)
+}

crates/desktop/src-tauri/src/utils/image/mod.rs

@@ -0,0 +1,3 @@
+pub mod decode_base64_image;
+pub mod encode_dynamic_image_base64_png;
+pub mod encode_rgb_image_base64_png;

crates/desktop/src-tauri/src/utils/mod.rs

@@ -0,0 +1 @@
+pub mod image;