From 7f7987b6e840b7c670be1ceea5b3e3e055ed3e76 Mon Sep 17 00:00:00 2001 From: Diretnan Domnan Date: Wed, 29 Jan 2025 16:14:09 +0100 Subject: [PATCH] =?UTF-8?q?Slight=20Resizing=20Optimizations=20=20for=20`p?= =?UTF-8?q?reprocess=5Fimages`=F0=9F=8F=83=E2=80=8D=E2=99=82=EF=B8=8F=20?= =?UTF-8?q?=20(#190)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit * Removing ndarray from DB * Attempting to reduce number of clones during preprocessing * Adding tracing for more viz * Attempting speedup with fast_resize * Adding notes for providers --- README.md | 5 + ahnlich/Cargo.lock | 31 +++ ahnlich/Cargo.toml | 2 +- ahnlich/ai/Cargo.toml | 1 + ahnlich/ai/src/engine/ai/models.rs | 225 +++++++++--------- .../ai/providers/processors/center_crop.rs | 8 +- .../processors/imagearray_to_ndarray.rs | 18 +- .../ai/providers/processors/normalize.rs | 1 + .../ai/providers/processors/preprocessor.rs | 2 + .../engine/ai/providers/processors/rescale.rs | 1 + .../engine/ai/providers/processors/resize.rs | 1 + .../ai/providers/processors/tokenize.rs | 1 + ahnlich/ai/src/error.rs | 4 +- ahnlich/ai/src/manager/mod.rs | 2 +- 14 files changed, 168 insertions(+), 134 deletions(-) diff --git a/README.md b/README.md index 1350b39a..52adb01c 100644 --- a/README.md +++ b/README.md @@ -170,6 +170,11 @@ services: ``` +### Execution Providers (Ahnlich AI) + +`CUDA`: Only supports >= CUDAv12 and might need to `sudo apt install libcudnn9-dev-cuda-12` +`CoreML (Apple)`: Not advised for NLP models due to often large dimensionality. + ### Contributing View [contribution guide](CONTRIBUTING.md) diff --git a/ahnlich/Cargo.lock b/ahnlich/Cargo.lock index e382f86e..628491df 100644 --- a/ahnlich/Cargo.lock +++ b/ahnlich/Cargo.lock @@ -95,6 +95,7 @@ dependencies = [ "deadpool", "dirs", "fallible_collections", + "fast_image_resize", "flurry", "futures", "hf-hub", @@ -1040,6 +1041,15 @@ dependencies = [ "syn 2.0.87", ] +[[package]] +name = "document-features" +version = "0.2.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cb6969eaabd2421f8a2775cfd2471a2b634372b4a25d41e3bd647b79912850a0" +dependencies = [ + "litrs", +] + [[package]] name = "dsl" version = "0.1.0" @@ -1157,6 +1167,21 @@ dependencies = [ "regex", ] +[[package]] +name = "fast_image_resize" +version = "5.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3d6b088992b0c2db53860aa4b43f299fd45aa85b835c744dec53c3f40231330d" +dependencies = [ + "bytemuck", + "cfg-if", + "document-features", + "image", + "num-traits", + "rayon", + "thiserror", +] + [[package]] name = "fastrand" version = "1.9.0" @@ -1961,6 +1986,12 @@ version = "0.7.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "643cb0b8d4fcc284004d5fd0d67ccf61dfffadb7f75e1e71bc420f4688a3a704" +[[package]] +name = "litrs" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b4ce301924b7887e9d637144fdade93f9dfff9b60981d4ac161db09720d39aa5" + [[package]] name = "lock_api" version = "0.4.12" diff --git a/ahnlich/Cargo.toml b/ahnlich/Cargo.toml index 51db6374..9268c391 100644 --- a/ahnlich/Cargo.toml +++ b/ahnlich/Cargo.toml @@ -19,7 +19,7 @@ async-trait = "0.1" serde = { version = "1.0.*", features = ["derive", "rc"] } bincode = "1.3.3" ndarray = { version = "0.16.1", features = ["serde", "rayon"] } -image = "0.25.2" +image = "0.25.5" serde_json = "1.0.116" itertools = "0.10.0" clap = { version = "4.5.4", features = ["derive"] } diff --git a/ahnlich/ai/Cargo.toml b/ahnlich/ai/Cargo.toml index 95ca8a87..6a6cac69 100644 --- a/ahnlich/ai/Cargo.toml +++ b/ahnlich/ai/Cargo.toml @@ -43,6 +43,7 @@ dirs = "5.0.1" ort = { version = "=2.0.0-rc.5", features = [ "ndarray", ] } +fast_image_resize = { version = "5.1.1", features = ["rayon"]} ort-sys = "=2.0.0-rc.8" moka = { version = "0.12.8", features = ["future", "sync"] } tracing-opentelemetry.workspace = true diff --git a/ahnlich/ai/src/engine/ai/models.rs b/ahnlich/ai/src/engine/ai/models.rs index 1a18b5ce..91264ea4 100644 --- a/ahnlich/ai/src/engine/ai/models.rs +++ b/ahnlich/ai/src/engine/ai/models.rs @@ -5,13 +5,21 @@ use crate::engine::ai::providers::ProviderTrait; use crate::error::AIProxyError; use ahnlich_types::ai::ExecutionProvider; use ahnlich_types::{ai::AIStoreInputType, keyval::StoreKey}; -use image::{DynamicImage, GenericImageView, ImageFormat, ImageReader}; +use fast_image_resize::images::Image; +use fast_image_resize::images::ImageRef; +use fast_image_resize::FilterType; +use fast_image_resize::PixelType; +use fast_image_resize::ResizeAlg; +use fast_image_resize::ResizeOptions; +use fast_image_resize::Resizer; +use image::imageops; +use image::ImageReader; +use image::RgbImage; use ndarray::{Array, Ix3}; use ndarray::{ArrayView, Ix4}; use nonzero_ext::nonzero; -use serde::de::Error as DeError; -use serde::ser::Error as SerError; -use serde::{Deserialize, Deserializer, Serialize, Serializer}; +use once_cell::sync::Lazy; +use serde::{Deserialize, Serialize}; use std::fmt; use std::io::Cursor; use std::num::NonZeroUsize; @@ -19,6 +27,8 @@ use std::path::PathBuf; use strum::Display; use tokenizers::Encoding; +static CHANNELS: Lazy = Lazy::new(|| image::ColorType::Rgb8.channel_count()); + #[derive(Display, Debug, Serialize, Deserialize)] pub enum ModelType { Text { @@ -246,31 +256,71 @@ pub enum ModelInput { Images(Array), } -#[derive(Debug, Clone)] -pub struct ImageArray { +#[derive(Debug)] +pub struct OnnxTransformResult { array: Array, - image: DynamicImage, - image_format: ImageFormat, - onnx_transformed: bool, } -impl ImageArray { - pub fn try_new(bytes: Vec) -> Result { - let img_reader = ImageReader::new(Cursor::new(&bytes)) - .with_guessed_format() - .map_err(|_| AIProxyError::ImageBytesDecodeError)?; +impl OnnxTransformResult { + pub fn view(&self) -> ArrayView { + self.array.view() + } - let image_format = &img_reader - .format() - .ok_or(AIProxyError::ImageBytesDecodeError)?; + pub fn image_dim(&self) -> (NonZeroUsize, NonZeroUsize) { + let shape = self.array.shape(); + ( + NonZeroUsize::new(shape[2]).expect("Array columns should be non zero"), + NonZeroUsize::new(shape[1]).expect("Array channels should be non zero"), + ) + } +} - let image = img_reader - .decode() +impl TryFrom for OnnxTransformResult { + type Error = AIProxyError; + + // Swapping axes from [rows, columns, channels] to [channels, rows, columns] for ONNX + #[tracing::instrument(skip_all)] + fn try_from(value: ImageArray) -> Result { + let image = value.image; + let mut array = Array::from_shape_vec( + ( + image.height() as usize, + image.width() as usize, + *CHANNELS as usize, + ), + image.into_raw(), + ) + .map_err(|e| AIProxyError::ImageArrayToNdArrayError { + message: format!("Error running onnx transform {e}"), + })? + .mapv(f32::from); + array.swap_axes(1, 2); + array.swap_axes(0, 1); + Ok(Self { array }) + } +} + +#[derive(Debug)] +pub struct ImageArray { + image: RgbImage, +} + +impl TryFrom<&[u8]> for ImageArray { + type Error = AIProxyError; + + #[tracing::instrument(skip_all)] + fn try_from(value: &[u8]) -> Result { + let img_reader = ImageReader::new(Cursor::new(value)) + .with_guessed_format() .map_err(|_| AIProxyError::ImageBytesDecodeError)?; // Always convert to RGB8 format // https://github.com/Anush008/fastembed-rs/blob/cea92b6c8b877efda762393848d1c449a4eea126/src/image_embedding/utils.rs#L198 - let image: DynamicImage = image.to_owned().into_rgb8().into(); + let image = img_reader + .decode() + .map_err(|_| AIProxyError::ImageBytesDecodeError)? + .into_rgb8(); + let (width, height) = image.dimensions(); if width == 0 || height == 0 { @@ -279,116 +329,59 @@ impl ImageArray { height: height as usize, }); } - - let channels = &image.color().channel_count(); - let shape = (height as usize, width as usize, *channels as usize); - let array = Array::from_shape_vec(shape, image.clone().into_bytes()) - .map_err(|_| AIProxyError::ImageBytesDecodeError)? - .mapv(f32::from); - - Ok(ImageArray { - array, - image, - image_format: image_format.to_owned(), - onnx_transformed: false, - }) - } - - // Swapping axes from [rows, columns, channels] to [channels, rows, columns] for ONNX - pub fn onnx_transform(&mut self) { - if self.onnx_transformed { - return; - } - self.array.swap_axes(1, 2); - self.array.swap_axes(0, 1); - self.onnx_transformed = true; - } - - pub fn view(&self) -> ArrayView { - self.array.view() + Ok(Self { image }) } +} - pub fn get_bytes(&self) -> Result, AIProxyError> { - let mut buffer = Cursor::new(Vec::new()); - let _ = &self - .image - .write_to(&mut buffer, self.image_format) - .map_err(|_| AIProxyError::ImageBytesEncodeError)?; - let bytes = buffer.into_inner(); - Ok(bytes) +impl ImageArray { + fn array_view(&self) -> ArrayView { + let shape = ( + self.image.height() as usize, + self.image.width() as usize, + *CHANNELS as usize, + ); + let raw_bytes = self.image.as_raw(); + ArrayView::from_shape(shape, raw_bytes).expect("Image bytes decode error") } + #[tracing::instrument(skip(self))] pub fn resize( - &self, + &mut self, width: u32, height: u32, filter: Option, ) -> Result { - let filter_type = filter.unwrap_or(image::imageops::FilterType::CatmullRom); - let resized_img = self.image.resize_exact(width, height, filter_type); - let channels = resized_img.color().channel_count(); - let shape = (height as usize, width as usize, channels as usize); - - let flattened_pixels = resized_img.clone().into_bytes(); - let array = Array::from_shape_vec(shape, flattened_pixels) - .map_err(|_| AIProxyError::ImageResizeError)? - .mapv(f32::from); - Ok(ImageArray { - array, - image: resized_img, - image_format: self.image_format, - onnx_transformed: false, - }) + // Create container for data of destination image + let (width, height) = self.image.dimensions(); + let mut dest_image = Image::new(width, height, PixelType::U8x3); + let mut resizer = Resizer::new(); + resizer + .resize( + &ImageRef::new(width, height, self.image.as_raw(), PixelType::U8x3) + .map_err(|e| AIProxyError::ImageResizeError(e.to_string()))?, + &mut dest_image, + &ResizeOptions::new().resize_alg(ResizeAlg::Convolution(FilterType::CatmullRom)), + ) + .map_err(|e| AIProxyError::ImageResizeError(e.to_string()))?; + let resized_img = RgbImage::from_raw(width, height, dest_image.into_vec()) + .expect("Could not get image after resizing"); + Ok(ImageArray { image: resized_img }) } - pub fn crop(&self, x: u32, y: u32, width: u32, height: u32) -> Result { - let cropped_img = self.image.crop_imm(x, y, width, height); - let channels = cropped_img.color().channel_count(); - let shape = (height as usize, width as usize, channels as usize); - - let flattened_pixels = cropped_img.clone().into_bytes(); - let array = Array::from_shape_vec(shape, flattened_pixels) - .map_err(|_| AIProxyError::ImageCropError)? - .mapv(f32::from); - Ok(ImageArray { - array, - image: cropped_img, - image_format: self.image_format, - onnx_transformed: false, - }) - } - - pub fn image_dim(&self) -> (NonZeroUsize, NonZeroUsize) { - let shape = self.array.shape(); - match self.onnx_transformed { - true => ( - NonZeroUsize::new(shape[2]).expect("Array columns should be non-zero"), - NonZeroUsize::new(shape[1]).expect("Array channels should be non-zero"), - ), // (width, channels) - false => ( - NonZeroUsize::new(shape[1]).expect("Array columns should be non-zero"), - NonZeroUsize::new(shape[0]).expect("Array rows should be non-zero"), - ), // (width, height) - } - } -} + #[tracing::instrument(skip(self))] + pub fn crop(&mut self, x: u32, y: u32, width: u32, height: u32) -> Result { + let cropped_img = imageops::crop(&mut self.image, x, y, width, height).to_image(); -impl Serialize for ImageArray { - fn serialize(&self, serializer: S) -> Result - where - S: Serializer, - { - serializer.serialize_bytes(&self.get_bytes().map_err(S::Error::custom)?) + Ok(ImageArray { image: cropped_img }) } -} -impl<'de> Deserialize<'de> for ImageArray { - fn deserialize(deserializer: D) -> Result - where - D: Deserializer<'de>, - { - let bytes: Vec = Deserialize::deserialize(deserializer)?; - ImageArray::try_new(bytes).map_err(D::Error::custom) + pub fn image_dim(&self) -> (NonZeroUsize, NonZeroUsize) { + let arr_view = self.array_view(); + let shape = arr_view.shape(); + ( + NonZeroUsize::new(shape[1]).expect("Array columns should be non-zero"), + NonZeroUsize::new(shape[0]).expect("Array rows should be non-zero"), + ) } } diff --git a/ahnlich/ai/src/engine/ai/providers/processors/center_crop.rs b/ahnlich/ai/src/engine/ai/providers/processors/center_crop.rs index ee2e331d..1ceb0943 100644 --- a/ahnlich/ai/src/engine/ai/providers/processors/center_crop.rs +++ b/ahnlich/ai/src/engine/ai/providers/processors/center_crop.rs @@ -3,7 +3,7 @@ use crate::engine::ai::providers::processors::{ Preprocessor, PreprocessorData, CONV_NEXT_FEATURE_EXTRACTOR_CENTER_CROP_THRESHOLD, }; use crate::error::AIProxyError; -use rayon::iter::{IntoParallelRefIterator, ParallelIterator}; +use rayon::iter::{IntoParallelIterator, ParallelIterator}; pub struct CenterCrop { crop_size: (u32, u32), // (width, height) @@ -98,18 +98,18 @@ impl CenterCrop { } impl Preprocessor for CenterCrop { + #[tracing::instrument(skip_all)] fn process(&self, data: PreprocessorData) -> Result { match data { PreprocessorData::ImageArray(image_array) => { let processed = image_array - .par_iter() - .map(|image| { + .into_par_iter() + .map(|mut image| { let (width, height) = image.image_dim(); let width = width.get() as u32; let height = height.get() as u32; let (crop_width, crop_height) = self.crop_size; if crop_width == width && crop_height == height { - let image = image.to_owned(); Ok(image) } else if crop_width <= width || crop_height <= height { let x = (width - crop_width) / 2; diff --git a/ahnlich/ai/src/engine/ai/providers/processors/imagearray_to_ndarray.rs b/ahnlich/ai/src/engine/ai/providers/processors/imagearray_to_ndarray.rs index 65b6b223..77ca591d 100644 --- a/ahnlich/ai/src/engine/ai/providers/processors/imagearray_to_ndarray.rs +++ b/ahnlich/ai/src/engine/ai/providers/processors/imagearray_to_ndarray.rs @@ -1,22 +1,20 @@ +use crate::engine::ai::models::OnnxTransformResult; use crate::engine::ai::providers::processors::{Preprocessor, PreprocessorData}; use crate::error::AIProxyError; pub struct ImageArrayToNdArray; impl Preprocessor for ImageArrayToNdArray { + #[tracing::instrument(skip_all)] fn process(&self, data: PreprocessorData) -> Result { match data { - PreprocessorData::ImageArray(mut arrays) => { + PreprocessorData::ImageArray(arrays) => { // Not using par_iter_mut here because it messes up the order of the images - // TODO: Figure out if it's more expensive to use par_iter_mut with enumerate or - // just keep doing it sequentially - let array_views = arrays - .iter_mut() - .map(|image_arr| { - image_arr.onnx_transform(); - image_arr.view() - }) - .collect::>(); + let arrays = arrays + .into_iter() + .map(OnnxTransformResult::try_from) + .collect::, _>>()?; + let array_views: Vec<_> = arrays.iter().map(|a| a.view()).collect(); let pixel_values_array = ndarray::stack(ndarray::Axis(0), &array_views).map_err(|_| { diff --git a/ahnlich/ai/src/engine/ai/providers/processors/normalize.rs b/ahnlich/ai/src/engine/ai/providers/processors/normalize.rs index 543dc640..c427fed2 100644 --- a/ahnlich/ai/src/engine/ai/providers/processors/normalize.rs +++ b/ahnlich/ai/src/engine/ai/providers/processors/normalize.rs @@ -39,6 +39,7 @@ impl ImageNormalize { } impl Preprocessor for ImageNormalize { + #[tracing::instrument(skip_all)] fn process(&self, data: PreprocessorData) -> Result { match data { PreprocessorData::NdArray3C(array) => { diff --git a/ahnlich/ai/src/engine/ai/providers/processors/preprocessor.rs b/ahnlich/ai/src/engine/ai/providers/processors/preprocessor.rs index a396fe68..0712b09f 100644 --- a/ahnlich/ai/src/engine/ai/providers/processors/preprocessor.rs +++ b/ahnlich/ai/src/engine/ai/providers/processors/preprocessor.rs @@ -53,6 +53,7 @@ impl ORTImagePreprocessor { }) } + #[tracing::instrument(skip_all)] pub fn process(&self, data: Vec) -> Result, AIProxyError> { let mut data = PreprocessorData::ImageArray(data); data = match self.resize { @@ -145,6 +146,7 @@ impl ORTTextPreprocessor { }) } + #[tracing::instrument(skip(self, data))] pub fn process( &self, data: Vec, diff --git a/ahnlich/ai/src/engine/ai/providers/processors/rescale.rs b/ahnlich/ai/src/engine/ai/providers/processors/rescale.rs index 677a536f..7ca02b02 100644 --- a/ahnlich/ai/src/engine/ai/providers/processors/rescale.rs +++ b/ahnlich/ai/src/engine/ai/providers/processors/rescale.rs @@ -18,6 +18,7 @@ impl Rescale { } impl Preprocessor for Rescale { + #[tracing::instrument(skip_all)] fn process(&self, data: PreprocessorData) -> Result { match data { PreprocessorData::NdArray3C(array) => { diff --git a/ahnlich/ai/src/engine/ai/providers/processors/resize.rs b/ahnlich/ai/src/engine/ai/providers/processors/resize.rs index ea7d743a..b914559a 100644 --- a/ahnlich/ai/src/engine/ai/providers/processors/resize.rs +++ b/ahnlich/ai/src/engine/ai/providers/processors/resize.rs @@ -94,6 +94,7 @@ impl Resize { } impl Preprocessor for Resize { + #[tracing::instrument(skip_all)] fn process(&self, data: PreprocessorData) -> Result { match data { PreprocessorData::ImageArray(mut arrays) => { diff --git a/ahnlich/ai/src/engine/ai/providers/processors/tokenize.rs b/ahnlich/ai/src/engine/ai/providers/processors/tokenize.rs index dec15640..df86c98d 100644 --- a/ahnlich/ai/src/engine/ai/providers/processors/tokenize.rs +++ b/ahnlich/ai/src/engine/ai/providers/processors/tokenize.rs @@ -139,6 +139,7 @@ impl Tokenize { } impl Preprocessor for Tokenize { + #[tracing::instrument(skip(self, data))] fn process(&self, data: PreprocessorData) -> Result { match data { PreprocessorData::Text(text) => { diff --git a/ahnlich/ai/src/error.rs b/ahnlich/ai/src/error.rs index 5a58df6b..26031a1d 100644 --- a/ahnlich/ai/src/error.rs +++ b/ahnlich/ai/src/error.rs @@ -127,8 +127,8 @@ pub enum AIProxyError { )] ImageNonzeroDimensionError { width: usize, height: usize }, - #[error("Image could not be resized.")] - ImageResizeError, + #[error("Image could not be resized. {0}")] + ImageResizeError(String), #[error("Image could not be cropped.")] ImageCropError, diff --git a/ahnlich/ai/src/manager/mod.rs b/ahnlich/ai/src/manager/mod.rs index 87236186..5ec3f154 100644 --- a/ahnlich/ai/src/manager/mod.rs +++ b/ahnlich/ai/src/manager/mod.rs @@ -108,7 +108,7 @@ impl ModelThread { .into_par_iter() .filter_map(|input| match input { StoreInput::Image(image_bytes) => { - Some(ImageArray::try_new(image_bytes).ok()?) + Some(ImageArray::try_from(image_bytes.as_slice()).ok()?) } _ => None, })