C++ and library part of iamge embeddings. Docs and examples will follow

Author: pweglik

packages/react-native-executorch/common/rnexecutorch/RnExecutorchInstaller.cpp

@@ -2,6 +2,7 @@
 
 #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>
@@ -42,6 +43,10 @@ void RnExecutorchInstaller::injectJSIBindings(
       *jsiRuntime, "loadExecutorchModule",
       RnExecutorchInstaller::loadModel<BaseModel>(jsiRuntime, jsCallInvoker,
                                                   "loadExecutorchModule"));
+  jsiRuntime->global().setProperty(
+      *jsiRuntime, "loadImageEmbeddings",
+      RnExecutorchInstaller::loadModel<ImageEmbeddings>(
+          jsiRuntime, jsCallInvoker, "loadImageEmbeddings"));
 }
 
 } // namespace rnexecutorch

packages/react-native-executorch/common/rnexecutorch/host_objects/JsiConversions.h

@@ -182,6 +182,23 @@ getJsiValue(const std::vector<std::shared_ptr<OwningArrayBuffer>> &vec,
   return jsi::Value(runtime, array);
 }
 
+inline jsi::Value getJsiValue(const std::shared_ptr<JSTensorViewOut> &jsTensor,
+                              jsi::Runtime &runtime) {
+
+  jsi::Object tensorObj(runtime);
+
+  tensorObj.setProperty(runtime, "sizes",
+                        getJsiValue(jsTensor->sizes, runtime));
+
+  tensorObj.setProperty(runtime, "scalarType",
+                        jsi::Value(static_cast<int>(jsTensor->scalarType)));
+
+  jsi::ArrayBuffer arrayBuffer(runtime, jsTensor->dataPtr);
+  tensorObj.setProperty(runtime, "dataPtr", arrayBuffer);
+
+  return jsi::Value(runtime, tensorObj);
+}
+
 inline jsi::Value
 getJsiValue(const std::vector<std::shared_ptr<JSTensorViewOut>> &vec,
             jsi::Runtime &runtime) {

packages/react-native-executorch/common/rnexecutorch/models/image_embeddings/ImageEmbeddings.cpp

@@ -0,0 +1,61 @@
+#include "ImageEmbeddings.h"
+
+#include <cstdint>
+#include <executorch/extension/tensor/tensor.h>
+#include <iostream>
+#include <rnexecutorch/Log.h>
+#include <rnexecutorch/data_processing/ImageProcessing.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]);
+}
+
+std::shared_ptr<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();
+  std::vector<int32_t> 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::memcpy(buffer->data(), outputTensor.const_data_ptr(), bufferSize);
+  auto jsTensor = std::make_shared<JSTensorViewOut>(
+      sizes, outputTensor.scalar_type(), buffer);
+
+  return jsTensor;
+}
+} // namespace rnexecutorch

packages/react-native-executorch/common/rnexecutorch/models/image_embeddings/ImageEmbeddings.h

@@ -0,0 +1,25 @@
+#pragma once
+
+#include <unordered_map>
+
+#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);
+  std::shared_ptr<JSTensorViewOut> generate(std::string imageSource);
+
+private:
+  cv::Size modelImageSize{0, 0};
+};
+
+} // namespace rnexecutorch

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,
+  });

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,

packages/react-native-executorch/src/index.tsx

@@ -9,6 +9,7 @@ declare global {
   var loadClassification: (source: string) => any;
   var loadObjectDetection: (source: string) => any;
   var loadExecutorchModule: (source: string) => any;
+  var loadImageEmbeddings: (source: string) => any;
 }
 // eslint-disable no-var
 if (
@@ -32,6 +33,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';
@@ -48,6 +50,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';

packages/react-native-executorch/src/modules/computer_vision/ImageEmbeddingsModule.ts

@@ -0,0 +1,24 @@
+import { ResourceFetcher } from '../../utils/ResourceFetcher';
+import { 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<TensorPtr> {
+    if (this.nativeModule == null)
+      throw new Error(getError(ETError.ModuleNotLoaded));
+    return await this.nativeModule.generate(imageSource);
+  }
+}