feat: allow embedding image bytes, Nix environment (#150)

* allow embedding images from their byte representation

* address feedback

---------

Co-authored-by: Stephen D <[email protected]>

Author: scd31

.envrc

@@ -0,0 +1 @@
+use flake

.gitignore

@@ -75,3 +75,6 @@ tests/rustdoc-gui/src/**.lock
 ## Rust files
 main.rs
 Cargo.lock
+
+## Nix
+/.direnv

flake.lock

@@ -0,0 +1,20 @@
+{
+  inputs = {
+    nixpkgs.url = "github:NixOS/nixpkgs/nixos-unstable";
+    flake-utils.url = "github:numtide/flake-utils?ref=main";
+  };
+
+  outputs = inputs:
+    inputs.flake-utils.lib.eachDefaultSystem (system:
+      let
+        pkgs = inputs.nixpkgs.legacyPackages.${system};
+
+      in {
+        devShells.default = pkgs.mkShell {
+          packages = (with pkgs; [
+            openssl
+            pkg-config
+          ]);
+        };
+      });
+}

flake.nix

@@ -3,14 +3,15 @@ use hf_hub::{
     api::sync::{ApiBuilder, ApiRepo},
     Cache,
 };
+use image::DynamicImage;
 use ndarray::{Array3, ArrayView3};
 use ort::{
     session::{builder::GraphOptimizationLevel, Session},
     value::Value,
 };
 #[cfg(feature = "hf-hub")]
 use std::path::PathBuf;
-use std::{path::Path, thread::available_parallelism};
+use std::{io::Cursor, path::Path, thread::available_parallelism};
 
 use crate::{
     common::normalize, models::image_embedding::models_list, Embedding, ImageEmbeddingModel,
@@ -132,14 +133,12 @@ impl ImageEmbedding {
             .expect("Model not found.")
     }
 
-    /// Method to generate image embeddings for a Vec of image path
-    // Generic type to accept String, &str, OsString, &OsStr
-    pub fn embed<S: AsRef<Path> + Send + Sync>(
+    /// Method to generate image embeddings for a Vec of image bytes
+    pub fn embed_bytes(
         &self,
-        images: Vec<S>,
+        images: &[&[u8]],
         batch_size: Option<usize>,
     ) -> anyhow::Result<Vec<Embedding>> {
-        // Determine the batch size, default if not specified
         let batch_size = batch_size.unwrap_or(DEFAULT_BATCH_SIZE);
 
         let output = images
@@ -149,72 +148,47 @@ impl ImageEmbedding {
                 let inputs = batch
                     .iter()
                     .map(|img| {
-                        let img = image::ImageReader::open(img)?
+                        image::ImageReader::new(Cursor::new(img))
+                            .with_guessed_format()?
                             .decode()
-                            .map_err(|err| anyhow!("image decode: {}", err))?;
-                        let pixels = self.preprocessor.transform(TransformData::Image(img))?;
-                        match pixels {
-                            TransformData::NdArray(array) => Ok(array),
-                            _ => Err(anyhow!("Preprocessor configuration error!")),
-                        }
+                            .map_err(|err| anyhow!("image decode: {}", err))
                     })
-                    .collect::<anyhow::Result<Vec<Array3<f32>>>>()?;
-
-                // Extract the batch size
-                let inputs_view: Vec<ArrayView3<f32>> =
-                    inputs.iter().map(|img| img.view()).collect();
-                let pixel_values_array = ndarray::stack(ndarray::Axis(0), &inputs_view)?;
+                    .collect::<Result<_, _>>()?;
 
-                let input_name = self.session.inputs[0].name.clone();
-                let session_inputs = ort::inputs![
-                    input_name => Value::from_array(pixel_values_array)?,
-                ]?;
+                self.embed_images(inputs)
+            })
+            .collect::<anyhow::Result<Vec<_>>>()?
+            .into_iter()
+            .flatten()
+            .collect();
 
-                let outputs = self.session.run(session_inputs)?;
+        Ok(output)
+    }
 
-                // Try to get the only output key
-                // If multiple, then default to few known keys `image_embeds` and `last_hidden_state`
-                let last_hidden_state_key = match outputs.len() {
-                    1 => vec![outputs.keys().next().unwrap()],
-                    _ => vec!["image_embeds", "last_hidden_state"],
-                };
+    /// Method to generate image embeddings for a Vec of image path
+    // Generic type to accept String, &str, OsString, &OsStr
+    pub fn embed<S: AsRef<Path> + Send + Sync>(
+        &self,
+        images: Vec<S>,
+        batch_size: Option<usize>,
+    ) -> anyhow::Result<Vec<Embedding>> {
+        // Determine the batch size, default if not specified
+        let batch_size = batch_size.unwrap_or(DEFAULT_BATCH_SIZE);
 
-                // Extract tensor and handle different dimensionalities
-                let output_data = last_hidden_state_key
+        let output = images
+            .par_chunks(batch_size)
+            .map(|batch| {
+                // Encode the texts in the batch
+                let inputs = batch
                     .iter()
-                    .find_map(|&key| {
-                        outputs
-                            .get(key)
-                            .and_then(|v| v.try_extract_tensor::<f32>().ok())
+                    .map(|img| {
+                        image::ImageReader::open(img)?
+                            .decode()
+                            .map_err(|err| anyhow!("image decode: {}", err))
                     })
-                    .ok_or_else(|| anyhow!("Could not extract tensor from any known output key"))?;
-                let shape = output_data.shape();
-
-                let embeddings: Vec<Vec<f32>> = match shape.len() {
-                    3 => {
-                        // For 3D output [batch_size, sequence_length, hidden_size]
-                        // Take only the first token, sequence_length[0] (CLS token), embedding
-                        // and return [batch_size, hidden_size]
-                        (0..shape[0])
-                            .map(|batch_idx| {
-                                let cls_embedding =
-                                    output_data.slice(ndarray::s![batch_idx, 0, ..]).to_vec();
-                                normalize(&cls_embedding)
-                            })
-                            .collect()
-                    }
-                    2 => {
-                        // For 2D output [batch_size, hidden_size]
-                        output_data
-                            .rows()
-                            .into_iter()
-                            .map(|row| normalize(row.as_slice().unwrap()))
-                            .collect()
-                    }
-                    _ => return Err(anyhow!("Unexpected output tensor shape: {:?}", shape)),
-                };
-
-                Ok(embeddings)
+                    .collect::<Result<_, _>>()?;
+
+                self.embed_images(inputs)
             })
             .collect::<anyhow::Result<Vec<_>>>()?
             .into_iter()
@@ -223,4 +197,73 @@ impl ImageEmbedding {
 
         Ok(output)
     }
+
+    /// Embed DynamicImages
+    pub fn embed_images(&self, imgs: Vec<DynamicImage>) -> anyhow::Result<Vec<Embedding>> {
+        let inputs = imgs
+            .into_iter()
+            .map(|img| {
+                let pixels = self.preprocessor.transform(TransformData::Image(img))?;
+                match pixels {
+                    TransformData::NdArray(array) => Ok(array),
+                    _ => Err(anyhow!("Preprocessor configuration error!")),
+                }
+            })
+            .collect::<anyhow::Result<Vec<Array3<f32>>>>()?;
+
+        // Extract the batch size
+        let inputs_view: Vec<ArrayView3<f32>> = inputs.iter().map(|img| img.view()).collect();
+        let pixel_values_array = ndarray::stack(ndarray::Axis(0), &inputs_view)?;
+
+        let input_name = self.session.inputs[0].name.clone();
+        let session_inputs = ort::inputs![
+            input_name => Value::from_array(pixel_values_array)?,
+        ]?;
+
+        let outputs = self.session.run(session_inputs)?;
+
+        // Try to get the only output key
+        // If multiple, then default to few known keys `image_embeds` and `last_hidden_state`
+        let last_hidden_state_key = match outputs.len() {
+            1 => vec![outputs.keys().next().unwrap()],
+            _ => vec!["image_embeds", "last_hidden_state"],
+        };
+
+        // Extract tensor and handle different dimensionalities
+        let output_data = last_hidden_state_key
+            .iter()
+            .find_map(|&key| {
+                outputs
+                    .get(key)
+                    .and_then(|v| v.try_extract_tensor::<f32>().ok())
+            })
+            .ok_or_else(|| anyhow!("Could not extract tensor from any known output key"))?;
+        let shape = output_data.shape();
+
+        let embeddings = match shape.len() {
+            3 => {
+                // For 3D output [batch_size, sequence_length, hidden_size]
+                // Take only the first token, sequence_length[0] (CLS token), embedding
+                // and return [batch_size, hidden_size]
+                (0..shape[0])
+                    .map(|batch_idx| {
+                        let cls_embedding =
+                            output_data.slice(ndarray::s![batch_idx, 0, ..]).to_vec();
+                        normalize(&cls_embedding)
+                    })
+                    .collect()
+            }
+            2 => {
+                // For 2D output [batch_size, hidden_size]
+                output_data
+                    .rows()
+                    .into_iter()
+                    .map(|row| normalize(row.as_slice().unwrap()))
+                    .collect()
+            }
+            _ => return Err(anyhow!("Unexpected output tensor shape: {:?}", shape)),
+        };
+
+        Ok(embeddings)
+    }
 }