C++ and library part of image embeddings.

pweglik · pweglik · commit 7808ff0e0563 · 2025-06-24T10:58:36.000+02:00
diff --git a/packages/react-native-executorch/common/rnexecutorch/RnExecutorchInstaller.cpp b/packages/react-native-executorch/common/rnexecutorch/RnExecutorchInstaller.cpp
@@ -3,6 +3,7 @@
 #include <rnexecutorch/TokenizerModule.h>
 #include <rnexecutorch/host_objects/JsiConversions.h>
 #include <rnexecutorch/models/classification/Classification.h>
+#include <rnexecutorch/models/image_embeddings/ImageEmbeddings.h>
 #include <rnexecutorch/models/image_segmentation/ImageSegmentation.h>
 #include <rnexecutorch/models/object_detection/ObjectDetection.h>
 #include <rnexecutorch/models/style_transfer/StyleTransfer.h>
@@ -48,6 +49,11 @@ void RnExecutorchInstaller::injectJSIBindings(
       *jsiRuntime, "loadTokenizerModule",
       RnExecutorchInstaller::loadModel<TokenizerModule>(
           jsiRuntime, jsCallInvoker, "loadTokenizerModule"));
+
+  jsiRuntime->global().setProperty(
+      *jsiRuntime, "loadImageEmbeddings",
+      RnExecutorchInstaller::loadModel<ImageEmbeddings>(
+          jsiRuntime, jsCallInvoker, "loadImageEmbeddings"));
 }
 
 } // namespace rnexecutorch
diff --git a/packages/react-native-executorch/common/rnexecutorch/RnExecutorchInstaller.h b/packages/react-native-executorch/common/rnexecutorch/RnExecutorchInstaller.h
@@ -30,6 +30,8 @@ REGISTER_CONSTRUCTOR(BaseModel, std::string,
                      std::shared_ptr<react::CallInvoker>);
 REGISTER_CONSTRUCTOR(TokenizerModule, std::string,
                      std::shared_ptr<react::CallInvoker>);
+REGISTER_CONSTRUCTOR(ImageEmbeddings, std::string,
+                     std::shared_ptr<react::CallInvoker>);
 
 using namespace facebook;
 
diff --git a/packages/react-native-executorch/common/rnexecutorch/TokenizerModule.cpp b/packages/react-native-executorch/common/rnexecutorch/TokenizerModule.cpp
@@ -1,7 +1,6 @@
 #include "TokenizerModule.h"
 #include <executorch/extension/module/module.h>
 #include <filesystem>
-#include <rnexecutorch/Log.h>
 #include <rnexecutorch/data_processing/FileUtils.h>
 
 namespace rnexecutorch {
diff --git a/packages/react-native-executorch/common/rnexecutorch/data_processing/Numerical.cpp b/packages/react-native-executorch/common/rnexecutorch/data_processing/Numerical.cpp
@@ -1,6 +1,7 @@
 #include "Numerical.h"
 
 #include <algorithm>
+#include <cmath>
 #include <numeric>
 
 namespace rnexecutorch::numerical {
@@ -16,4 +17,19 @@ void softmax(std::vector<float> &v) {
     x /= sum;
   }
 }
+void normalizeVector(std::vector<float> &v) {
+  float norm = 0.0;
+  for (float value : v) {
+    norm += value * value;
+  }
+  norm = sqrt(norm);
+
+  if (norm == 0) {
+    return;
+  }
+
+  for (float &value : v) {
+    value /= norm;
+  }
+}
 } // namespace rnexecutorch::numerical
diff --git a/packages/react-native-executorch/common/rnexecutorch/data_processing/Numerical.h b/packages/react-native-executorch/common/rnexecutorch/data_processing/Numerical.h
@@ -4,4 +4,5 @@
 
 namespace rnexecutorch::numerical {
 void softmax(std::vector<float> &v);
+void normalizeVector(std::vector<float> &v);
 } // namespace rnexecutorch::numerical
diff --git a/packages/react-native-executorch/common/rnexecutorch/host_objects/JsiConversions.h b/packages/react-native-executorch/common/rnexecutorch/host_objects/JsiConversions.h
@@ -216,6 +216,21 @@ inline jsi::Value getJsiValue(const std::vector<JSTensorViewOut> &vec,
   return jsi::Value(runtime, array);
 }
 
+inline jsi::Value getJsiValue(const JSTensorViewOut &tensor,
+                              jsi::Runtime &runtime) {
+  jsi::Object tensorObj(runtime);
+
+  tensorObj.setProperty(runtime, "sizes", getJsiValue(tensor.sizes, runtime));
+
+  tensorObj.setProperty(runtime, "scalarType",
+                        jsi::Value(static_cast<int>(tensor.scalarType)));
+
+  jsi::ArrayBuffer arrayBuffer(runtime, tensor.dataPtr);
+  tensorObj.setProperty(runtime, "dataPtr", arrayBuffer);
+
+  return tensorObj;
+}
+
 inline jsi::Value getJsiValue(const std::string &str, jsi::Runtime &runtime) {
   return jsi::String::createFromAscii(runtime, str);
 }
diff --git a/packages/react-native-executorch/common/rnexecutorch/models/image_embeddings/ImageEmbeddings.cpp b/packages/react-native-executorch/common/rnexecutorch/models/image_embeddings/ImageEmbeddings.cpp
@@ -0,0 +1,66 @@
+#include "ImageEmbeddings.h"
+
+#include <executorch/extension/tensor/tensor.h>
+#include <rnexecutorch/Log.h>
+#include <rnexecutorch/data_processing/ImageProcessing.h>
+#include <rnexecutorch/data_processing/Numerical.h>
+namespace rnexecutorch {
+
+ImageEmbeddings::ImageEmbeddings(
+    const std::string &modelSource,
+    std::shared_ptr<react::CallInvoker> callInvoker)
+    : BaseModel(modelSource, callInvoker) {
+  auto inputTensors = getAllInputShapes();
+  if (inputTensors.size() == 0) {
+    throw std::runtime_error("Model seems to not take any input tensors.");
+  }
+  std::vector<int32_t> modelInputShape = inputTensors[0];
+  if (modelInputShape.size() < 2) {
+    char errorMessage[100];
+    std::snprintf(errorMessage, sizeof(errorMessage),
+                  "Unexpected model input size, expected at least 2 dimentions "
+                  "but got: %zu.",
+                  modelInputShape.size());
+    throw std::runtime_error(errorMessage);
+  }
+  modelImageSize = cv::Size(modelInputShape[modelInputShape.size() - 1],
+                            modelInputShape[modelInputShape.size() - 2]);
+}
+
+JSTensorViewOut ImageEmbeddings::generate(std::string imageSource) {
+  auto [inputTensor, originalSize] =
+      imageprocessing::readImageToTensor(imageSource, getAllInputShapes()[0]);
+
+  auto result = BaseModel::forward(inputTensor);
+  if (!result.ok()) {
+    throw std::runtime_error("Forward pass failed: Error " +
+                             std::to_string(static_cast<int>(result.error())));
+  }
+
+  auto &outputs = result.get();
+
+  if (outputs.size() > 1) {
+    throw std::runtime_error("It returned multiple outputs!");
+  }
+
+  auto &outputTensor = outputs.at(0).toTensor();
+  auto sizesRaw = outputTensor.sizes();
+  auto sizes = std::vector<int32_t>(sizesRaw.begin(), sizesRaw.end());
+  size_t bufferSize = outputTensor.numel() * outputTensor.element_size();
+  auto buffer = std::make_shared<OwningArrayBuffer>(bufferSize);
+
+  std::span<const float> outputTensorSpan(
+      static_cast<const float *>(outputTensor.const_data_ptr()),
+      outputTensor.numel());
+  std::vector<float> outputVector(outputTensorSpan.begin(),
+                                  outputTensorSpan.end());
+
+  numerical::normalizeVector(outputVector);
+
+  std::memcpy(buffer->data(), outputVector.data(), bufferSize);
+
+  auto jsTensor = JSTensorViewOut(sizes, outputTensor.scalar_type(), buffer);
+
+  return jsTensor;
+}
+} // namespace rnexecutorch
diff --git a/packages/react-native-executorch/common/rnexecutorch/models/image_embeddings/ImageEmbeddings.h b/packages/react-native-executorch/common/rnexecutorch/models/image_embeddings/ImageEmbeddings.h
@@ -0,0 +1,23 @@
+#pragma once
+
+#include <executorch/extension/tensor/tensor_ptr.h>
+#include <executorch/runtime/core/evalue.h>
+#include <opencv2/opencv.hpp>
+
+#include <rnexecutorch/models/BaseModel.h>
+
+namespace rnexecutorch {
+using executorch::extension::TensorPtr;
+using executorch::runtime::EValue;
+
+class ImageEmbeddings : public BaseModel {
+public:
+  ImageEmbeddings(const std::string &modelSource,
+                  std::shared_ptr<react::CallInvoker> callInvoker);
+  JSTensorViewOut generate(std::string imageSource);
+
+private:
+  cv::Size modelImageSize{0, 0};
+};
+
+} // namespace rnexecutorch
diff --git a/packages/react-native-executorch/src/constants/modelUrls.ts b/packages/react-native-executorch/src/constants/modelUrls.ts
@@ -162,6 +162,10 @@ export const DETECTOR_CRAFT_320 =
 export const DEEPLAB_V3_RESNET50 =
   'https://huggingface.co/software-mansion/react-native-executorch-deeplab-v3/resolve/v0.4.0/xnnpack/deeplabV3_xnnpack_fp32.pte';
 
+// Image Embeddings
+export const CLIP_VIT_BASE_PATCH_32_IMAGE_ENCODER_MODEL =
+  'https://huggingface.co/software-mansion/react-native-executorch-clip-vit-base-patch32-image-encoder/resolve/v0.5.0/clip-vit-base-patch32-image-encoder-float32.pte';
+
 // Text Embeddings
 export const ALL_MINILM_L6_V2 =
   'https://huggingface.co/software-mansion/react-native-executorch-all-MiniLM-L6-v2/resolve/v0.4.0/all-MiniLM-L6-v2_xnnpack.pte';
@@ -183,6 +187,17 @@ export const MULTI_QA_MPNET_BASE_DOT_V1 =
 export const MULTI_QA_MPNET_BASE_DOT_V1_TOKENIZER =
   'https://huggingface.co/software-mansion/react-native-executorch-multi-qa-mpnet-base-dot-v1/resolve/v0.4.0/tokenizer.json';
 
+export const CLIP_VIT_BASE_PATCH_32_TEXT_ENCODER_MODEL =
+  'https://huggingface.co/software-mansion/react-native-executorch-clip-vit-base-patch32-text-encoder/resolve/main/clip-vit-base-patch32-text-encoder-float32.pte';
+export const CLIP_VIT_BASE_PATCH_32_TEXT_ENCODER_TOKENIZER =
+  'https://huggingface.co/software-mansion/react-native-executorch-clip-vit-base-patch32-text-encoder/resolve/main/tokenizer.json';
+
+export const CLIP_VIT_BASE_PATCH_32_TEXT_ENCODER = {
+  modelSource: CLIP_VIT_BASE_PATCH_32_TEXT_ENCODER_MODEL,
+  tokenizerSource: CLIP_VIT_BASE_PATCH_32_TEXT_ENCODER_TOKENIZER,
+  meanPooling: false,
+};
+
 // Backward compatibility
 export const LLAMA3_2_3B_URL = LLAMA3_2_3B;
 export const LLAMA3_2_3B_QLORA_URL = LLAMA3_2_3B_QLORA;
diff --git a/packages/react-native-executorch/src/hooks/computer_vision/useImageEmbeddings.ts b/packages/react-native-executorch/src/hooks/computer_vision/useImageEmbeddings.ts
@@ -0,0 +1,16 @@
+import { ImageEmbeddingsModule } from '../../modules/computer_vision/ImageEmbeddingsModule';
+import { ResourceSource } from '../../types/common';
+import { useNonStaticModule } from '../useNonStaticModule';
+
+export const useImageEmbeddings = ({
+  modelSource,
+  preventLoad = false,
+}: {
+  modelSource: ResourceSource;
+  preventLoad?: boolean;
+}) =>
+  useNonStaticModule({
+    module: ImageEmbeddingsModule,
+    loadArgs: [modelSource],
+    preventLoad,
+  });
diff --git a/packages/react-native-executorch/src/hooks/useNonStaticModule.ts b/packages/react-native-executorch/src/hooks/useNonStaticModule.ts
@@ -14,7 +14,7 @@ interface ModuleConstructor<M extends Module> {
 export const useNonStaticModule = <
   M extends Module,
   LoadArgs extends Parameters<M['load']>,
-  ForwardArgs extends any[],
+  ForwardArgs extends Parameters<M['forward']>,
   ForwardReturn extends Awaited<ReturnType<M['forward']>>,
 >({
   module,
diff --git a/packages/react-native-executorch/src/index.tsx b/packages/react-native-executorch/src/index.tsx
@@ -10,6 +10,7 @@ declare global {
   var loadObjectDetection: (source: string) => any;
   var loadExecutorchModule: (source: string) => any;
   var loadTokenizerModule: (source: string) => any;
+  var loadImageEmbeddings: (source: string) => any;
 }
 // eslint-disable no-var
 if (
@@ -35,6 +36,7 @@ export * from './hooks/computer_vision/useStyleTransfer';
 export * from './hooks/computer_vision/useImageSegmentation';
 export * from './hooks/computer_vision/useOCR';
 export * from './hooks/computer_vision/useVerticalOCR';
+export * from './hooks/computer_vision/useImageEmbeddings';
 
 export * from './hooks/natural_language_processing/useLLM';
 export * from './hooks/natural_language_processing/useSpeechToText';
@@ -51,6 +53,7 @@ export * from './modules/computer_vision/ImageSegmentationModule';
 export * from './modules/computer_vision/OCRModule';
 export * from './modules/computer_vision/VerticalOCRModule';
 export * from './modules/general/ExecutorchModule';
+export * from './modules/computer_vision/ImageEmbeddingsModule';
 
 export * from './modules/natural_language_processing/LLMModule';
 export * from './modules/natural_language_processing/SpeechToTextModule';
diff --git a/packages/react-native-executorch/src/modules/computer_vision/ImageEmbeddingsModule.ts b/packages/react-native-executorch/src/modules/computer_vision/ImageEmbeddingsModule.ts
@@ -0,0 +1,28 @@
+import { ResourceFetcher } from '../../utils/ResourceFetcher';
+import { getArrayConstructor, ResourceSource } from '../../types/common';
+import { TensorPtr } from '../../types/common';
+import { ETError, getError } from '../../Error';
+import { BaseNonStaticModule } from '../BaseNonStaticModule';
+
+export class ImageEmbeddingsModule extends BaseNonStaticModule {
+  async load(
+    modelSource: ResourceSource,
+    onDownloadProgressCallback: (_: number) => void = () => {}
+  ): Promise<void> {
+    const paths = await ResourceFetcher.fetchMultipleResources(
+      onDownloadProgressCallback,
+      modelSource
+    );
+    this.nativeModule = global.loadImageEmbeddings(paths[0] || '');
+  }
+
+  async forward(imageSource: string): Promise<number[]> {
+    if (this.nativeModule == null)
+      throw new Error(getError(ETError.ModuleNotLoaded));
+    const tensor: TensorPtr = await this.nativeModule.generate(imageSource);
+
+    const resultArray = getArrayConstructor(tensor.scalarType)(tensor.dataPtr);
+
+    return Array.from(resultArray);
+  }
+}
diff --git a/packages/react-native-executorch/src/types/common.ts b/packages/react-native-executorch/src/types/common.ts
@@ -7,6 +7,16 @@ export const getTypeIdentifier = (input: ETInput): number => {
   return -1;
 };
 
+export const getArrayConstructor = (
+  scalarType: ScalarType
+): ((arr: TensorBuffer) => Float32Array) => {
+  if (scalarType === ScalarType.FLOAT)
+    return (arr: TensorBuffer) => new Float32Array(arr as Float32Array);
+  // todo add more types?
+
+  throw Error("We don't handle this type of output array");
+};
+
 export type ResourceSource = string | number | object;
 
 export type ETInput =
