fix: adapt segmentation for refactored promises

JakubGonera · JakubGonera · commit 353d8e617bb9 · 2025-05-21T16:02:30.000+02:00
diff --git a/common/rnexecutorch/RnExecutorchInstaller.h b/common/rnexecutorch/RnExecutorchInstaller.h
@@ -42,7 +42,7 @@ class RnExecutorchInstaller {
                 jsiconversion::getValue<std::string>(args[0], runtime);
 
             auto modelImplementationPtr =
-                std::make_shared<ModelT>(source, &runtime);
+                std::make_shared<ModelT>(source, jsCallInvoker);
             auto modelHostObject = std::make_shared<ModelHostObject<ModelT>>(
                 modelImplementationPtr, jsCallInvoker);
 
diff --git a/common/rnexecutorch/host_objects/JsiConversions.h b/common/rnexecutorch/host_objects/JsiConversions.h
@@ -71,8 +71,9 @@ getValue<std::set<std::string, std::less<>>>(const jsi::Value &val,
 // we add a function here.
 
 // Identity function for the sake of completeness
-inline jsi::Value getJsiValue(jsi::Value &&value, jsi::Runtime &runtime) {
-  return std::move(value);
+inline jsi::Value getJsiValue(std::unique_ptr<jsi::Object> &&valuePtr,
+                              jsi::Runtime &runtime) {
+  return std::move(*valuePtr);
 }
 
 inline jsi::Value getJsiValue(const std::string &str, jsi::Runtime &runtime) {
diff --git a/common/rnexecutorch/host_objects/ModelHostObject.h b/common/rnexecutorch/host_objects/ModelHostObject.h
@@ -51,8 +51,8 @@ template <typename Model> class ModelHostObject : public JsiHostObject {
                 auto result = std::apply(
                     std::bind_front(&Model::forward, model), argsConverted);
 
-                callInvoker->invokeAsync([promise, result = std::move(result)](
-                                             jsi::Runtime &runtime) {
+                callInvoker->invokeSync([promise,
+                                         &result](jsi::Runtime &runtime) {
                   promise->resolve(
                       jsiconversion::getJsiValue(std::move(result), runtime));
                 });
diff --git a/common/rnexecutorch/models/BaseModel.cpp b/common/rnexecutorch/models/BaseModel.cpp
@@ -4,14 +4,15 @@
 
 namespace rnexecutorch {
 
+using namespace facebook;
 using ::executorch::extension::Module;
 using ::executorch::runtime::Error;
 
 BaseModel::BaseModel(const std::string &modelSource,
-                     facebook::jsi::Runtime *runtime)
+                     std::shared_ptr<react::CallInvoker> callInvoker)
     : module(std::make_unique<Module>(
           modelSource, Module::LoadMode::MmapUseMlockIgnoreErrors)),
-      runtime(runtime) {
+      callInvoker(callInvoker) {
   Error loadError = module->load();
   if (loadError != Error::Ok) {
     throw std::runtime_error("Couldn't load the model, error: " +
diff --git a/common/rnexecutorch/models/BaseModel.h b/common/rnexecutorch/models/BaseModel.h
@@ -2,22 +2,25 @@
 
 #include <string>
 
+#include <ReactCommon/CallInvoker.h>
 #include <executorch/extension/module/module.h>
 #include <jsi/jsi.h>
 
 namespace rnexecutorch {
+using namespace facebook;
 
 class BaseModel {
 public:
-  BaseModel(const std::string &modelSource, facebook::jsi::Runtime *runtime);
+  BaseModel(const std::string &modelSource,
+            std::shared_ptr<react::CallInvoker> callInvoker);
   std::vector<std::vector<int32_t>> getInputShape();
 
 protected:
   std::unique_ptr<executorch::extension::Module> module;
-  // If possible, models should not use the runtime to keep JSI internals away
-  // from logic, however, sometimes this would incur too big of a penalty
-  // (unnecessary copies). This is in BaseModel so that we can generalize JSI
-  // loader method installation.
-  facebook::jsi::Runtime *runtime;
+  // If possible, models should not use the JS runtime to keep JSI internals
+  // away from logic, however, sometimes this would incur too big of a penalty
+  // (unnecessary copies instead of working on JS memory). In this case
+  // CallInvoker can be used to get jsi::Runtime, and use it in a safe manner.
+  std::shared_ptr<react::CallInvoker> callInvoker;
 };
 } // namespace rnexecutorch
diff --git a/common/rnexecutorch/models/image_segmentation/ImageSegmentation.cpp b/common/rnexecutorch/models/image_segmentation/ImageSegmentation.cpp
@@ -9,17 +9,18 @@
 
 namespace rnexecutorch {
 
-ImageSegmentation::ImageSegmentation(const std::string &modelSource,
-                                     jsi::Runtime *runtime)
-    : BaseModel(modelSource, runtime) {
+ImageSegmentation::ImageSegmentation(
+    const std::string &modelSource,
+    std::shared_ptr<react::CallInvoker> callInvoker)
+    : BaseModel(modelSource, callInvoker) {
 
-  std::vector<int32_t> modelInputShape = getInputShape();
+  std::vector<int32_t> modelInputShape = getInputShape()[0];
   modelImageSize = cv::Size(modelInputShape[modelInputShape.size() - 1],
                             modelInputShape[modelInputShape.size() - 2]);
   numModelPixels = modelImageSize.area();
 }
 
-jsi::Value
+std::unique_ptr<jsi::Object>
 ImageSegmentation::forward(std::string imageSource,
                            std::set<std::string, std::less<>> classesOfInterest,
                            bool resize) {
@@ -36,7 +37,20 @@ ImageSegmentation::forward(std::string imageSource,
                      classesOfInterest, resize);
 }
 
-jsi::Value ImageSegmentation::postprocess(
+std::pair<TensorPtr, cv::Size>
+ImageSegmentation::preprocess(const std::string &imageSource) {
+  cv::Mat input = imageprocessing::readImage(imageSource);
+  cv::Size inputSize = input.size();
+
+  cv::resize(input, input, modelImageSize);
+
+  std::vector<float> inputVector = imageprocessing::colorMatToVector(input);
+  return {
+      executorch::extension::make_tensor_ptr(getInputShape()[0], inputVector),
+      inputSize};
+}
+
+std::unique_ptr<jsi::Object> ImageSegmentation::postprocess(
     const Tensor &tensor, cv::Size originalSize,
     std::set<std::string, std::less<>> classesOfInterest, bool resize) {
 
@@ -84,11 +98,11 @@ jsi::Value ImageSegmentation::postprocess(
     reinterpret_cast<int32_t *>(argmax->data())[pixel] = maxInd;
   }
 
-  std::unordered_map<std::string_view, std::shared_ptr<OwningArrayBuffer>>
-      buffersToReturn;
+  auto buffersToReturn = std::make_shared<std::unordered_map<
+      std::string_view, std::shared_ptr<OwningArrayBuffer>>>();
   for (std::size_t cl = 0; cl < numClasses; ++cl) {
     if (classesOfInterest.contains(deeplabv3_resnet50_labels[cl])) {
-      buffersToReturn[deeplabv3_resnet50_labels[cl]] = resultClasses[cl];
+      (*buffersToReturn)[deeplabv3_resnet50_labels[cl]] = resultClasses[cl];
     }
   }
 
@@ -102,7 +116,7 @@ jsi::Value ImageSegmentation::postprocess(
     std::memcpy(argmax->data(), argmaxMat.data,
                 originalSize.area() * sizeof(int32_t));
 
-    for (auto &[label, arrayBuffer] : buffersToReturn) {
+    for (auto &[label, arrayBuffer] : *buffersToReturn) {
       cv::Mat classMat(modelImageSize, CV_32FC1, arrayBuffer->data());
       cv::resize(classMat, classMat, originalSize);
       arrayBuffer = std::make_shared<OwningArrayBuffer>(originalSize.area() *
@@ -114,53 +128,41 @@ jsi::Value ImageSegmentation::postprocess(
   return populateDictionary(argmax, buffersToReturn);
 }
 
-jsi::Value ImageSegmentation::populateDictionary(
+std::unique_ptr<jsi::Object> ImageSegmentation::populateDictionary(
     std::shared_ptr<OwningArrayBuffer> argmax,
-    std::unordered_map<std::string_view, std::shared_ptr<OwningArrayBuffer>>
+    std::shared_ptr<std::unordered_map<std::string_view,
+                                       std::shared_ptr<OwningArrayBuffer>>>
         classesToOutput) {
-  jsi::Object dict(*runtime);
-
-  auto argmaxArrayBuffer = jsi::ArrayBuffer(*runtime, argmax);
-
-  auto int32ArrayCtor =
-      runtime->global().getPropertyAsFunction(*runtime, "Int32Array");
-  auto int32Array =
-      int32ArrayCtor.callAsConstructor(*runtime, argmaxArrayBuffer)
-          .getObject(*runtime);
-  dict.setProperty(*runtime, "ARGMAX", int32Array);
-
-  std::size_t dictIndex = 1;
-  for (auto &[classLabel, owningBuffer] : classesToOutput) {
-    auto classArrayBuffer = jsi::ArrayBuffer(*runtime, owningBuffer);
-
-    auto float32ArrayCtor =
-        runtime->global().getPropertyAsFunction(*runtime, "Float32Array");
-    auto float32Array =
-        float32ArrayCtor.callAsConstructor(*runtime, classArrayBuffer)
-            .getObject(*runtime);
-
-    dict.setProperty(*runtime,
-                     jsi::String::createFromAscii(*runtime, classLabel.data()),
-                     float32Array);
-  }
-  return dict;
-}
-
-std::pair<TensorPtr, cv::Size>
-ImageSegmentation::preprocess(const std::string &imageSource) {
-  cv::Mat input = imageprocessing::readImage(imageSource);
-  cv::Size inputSize = input.size();
-
-  std::vector<int32_t> modelInputShape = getInputShape();
-  cv::Size modelImageSize =
-      cv::Size(modelInputShape[modelInputShape.size() - 1],
-               modelInputShape[modelInputShape.size() - 2]);
-
-  cv::resize(input, input, modelImageSize);
-
-  std::vector<float> inputVector = imageprocessing::colorMatToVector(input);
-  return {executorch::extension::make_tensor_ptr(modelInputShape, inputVector),
-          inputSize};
+  std::unique_ptr<jsi::Object> dictPtr;
+
+  callInvoker->invokeSync(
+      [argmax, classesToOutput, &dictPtr](jsi::Runtime &runtime) {
+        dictPtr = std::make_unique<jsi::Object>(runtime);
+        auto argmaxArrayBuffer = jsi::ArrayBuffer(runtime, argmax);
+
+        auto int32ArrayCtor =
+            runtime.global().getPropertyAsFunction(runtime, "Int32Array");
+        auto int32Array =
+            int32ArrayCtor.callAsConstructor(runtime, argmaxArrayBuffer)
+                .getObject(runtime);
+        dictPtr->setProperty(runtime, "ARGMAX", int32Array);
+
+        for (auto &[classLabel, owningBuffer] : *classesToOutput) {
+          auto classArrayBuffer = jsi::ArrayBuffer(runtime, owningBuffer);
+
+          auto float32ArrayCtor =
+              runtime.global().getPropertyAsFunction(runtime, "Float32Array");
+          auto float32Array =
+              float32ArrayCtor.callAsConstructor(runtime, classArrayBuffer)
+                  .getObject(runtime);
+
+          dictPtr->setProperty(
+              runtime, jsi::String::createFromAscii(runtime, classLabel.data()),
+              float32Array);
+        }
+      });
+
+  return dictPtr;
 }
 
 } // namespace rnexecutorch
diff --git a/common/rnexecutorch/models/image_segmentation/ImageSegmentation.h b/common/rnexecutorch/models/image_segmentation/ImageSegmentation.h
@@ -19,19 +19,22 @@ using executorch::extension::TensorPtr;
 
 class ImageSegmentation : public BaseModel {
 public:
-  ImageSegmentation(const std::string &modelSource, jsi::Runtime *runtime);
-  jsi::Value forward(std::string imageSource,
-                     std::set<std::string, std::less<>> classesOfInterest,
-                     bool resize);
+  ImageSegmentation(const std::string &modelSource,
+                    std::shared_ptr<react::CallInvoker> callInvoker);
+  std::unique_ptr<jsi::Object>
+  forward(std::string imageSource,
+          std::set<std::string, std::less<>> classesOfInterest, bool resize);
 
 private:
   std::pair<TensorPtr, cv::Size> preprocess(const std::string &imageSource);
-  jsi::Value postprocess(const Tensor &tensor, cv::Size originalSize,
-                         std::set<std::string, std::less<>> classesOfInterest,
-                         bool resize);
-  jsi::Value populateDictionary(
+  std::unique_ptr<jsi::Object>
+  postprocess(const Tensor &tensor, cv::Size originalSize,
+              std::set<std::string, std::less<>> classesOfInterest,
+              bool resize);
+  std::unique_ptr<jsi::Object> populateDictionary(
       std::shared_ptr<OwningArrayBuffer> argmax,
-      std::unordered_map<std::string_view, std::shared_ptr<OwningArrayBuffer>>
+      std::shared_ptr<std::unordered_map<std::string_view,
+                                         std::shared_ptr<OwningArrayBuffer>>>
           classesToOutput);
 
   static constexpr std::size_t numClasses{deeplabv3_resnet50_labels.size()};
diff --git a/common/rnexecutorch/models/style_transfer/StyleTransfer.cpp b/common/rnexecutorch/models/style_transfer/StyleTransfer.cpp
@@ -15,8 +15,8 @@ using executorch::extension::TensorPtr;
 using executorch::runtime::Error;
 
 StyleTransfer::StyleTransfer(const std::string &modelSource,
-                             jsi::Runtime *runtime)
-    : BaseModel(modelSource, runtime) {
+                             std::shared_ptr<react::CallInvoker> callInvoker)
+    : BaseModel(modelSource, callInvoker) {
   std::vector<int32_t> modelInputShape = getInputShape()[0];
   modelImageSize = cv::Size(modelInputShape[modelInputShape.size() - 1],
                             modelInputShape[modelInputShape.size() - 2]);
diff --git a/common/rnexecutorch/models/style_transfer/StyleTransfer.h b/common/rnexecutorch/models/style_transfer/StyleTransfer.h
@@ -17,7 +17,8 @@ using executorch::extension::TensorPtr;
 
 class StyleTransfer : public BaseModel {
 public:
-  StyleTransfer(const std::string &modelSource, jsi::Runtime *runtime);
+  StyleTransfer(const std::string &modelSource,
+                std::shared_ptr<react::CallInvoker> callInvoker);
   std::string forward(std::string imageSource);
 
 private:
