feat: port image segmentation to C++

Author: JakubGonera

android/src/main/java/com/swmansion/rnexecutorch/ImageSegmentation.kt

@@ -28,8 +28,6 @@ class RnExecutorchPackage : TurboReactPackage() {
       OCR(reactContext)
     } else if (name == VerticalOCR.NAME) {
       VerticalOCR(reactContext)
-    } else if (name == ImageSegmentation.NAME) {
-      ImageSegmentation(reactContext)
     } else if (name == ETInstaller.NAME) {
       ETInstaller(reactContext)
     } else if (name == Tokenizer.NAME) {
@@ -119,17 +117,6 @@ class RnExecutorchPackage : TurboReactPackage() {
           true,
         )
 
-      moduleInfos[ImageSegmentation.NAME] =
-        ReactModuleInfo(
-          ImageSegmentation.NAME,
-          ImageSegmentation.NAME,
-          false, // canOverrideExistingModule
-          false, // needsEagerInit
-          true, // hasConstants
-          false, // isCxxModule
-          true,
-        )
-
       moduleInfos[Tokenizer.NAME] =
         ReactModuleInfo(
           Tokenizer.NAME,

android/src/main/java/com/swmansion/rnexecutorch/RnExecutorchPackage.kt

@@ -1,44 +1,26 @@
 #include "RnExecutorchInstaller.h"
 
 #include <rnexecutorch/host_objects/JsiConversions.h>
-#include <rnexecutorch/host_objects/ModelHostObject.h>
-#include <rnexecutorch/jsi/JsiPromise.h>
-#include <rnexecutorch/models/StyleTransfer.h>
+#include <rnexecutorch/models/image_segmentation/ImageSegmentation.h>
+#include <rnexecutorch/models/style_transfer/StyleTransfer.h>
 
 namespace rnexecutorch {
 
 FetchUrlFunc_t fetchUrlFunc;
 
-jsi::Function RnExecutorchInstaller::loadStyleTransfer(
-    jsi::Runtime *jsiRuntime,
-    const std::shared_ptr<react::CallInvoker> &jsCallInvoker) {
-  return jsi::Function::createFromHostFunction(
-      *jsiRuntime, jsi::PropNameID::forAscii(*jsiRuntime, "loadStyleTransfer"),
-      0,
-      [jsCallInvoker](jsi::Runtime &runtime, const jsi::Value &thisValue,
-                      const jsi::Value *args, size_t count) -> jsi::Value {
-        assert(count == 1);
-        auto source = jsiconversion::getValue<std::string>(args[0], runtime);
-
-        auto styleTransferPtr =
-            std::make_shared<StyleTransfer>(source, &runtime);
-        auto styleTransferHostObject =
-            std::make_shared<ModelHostObject<StyleTransfer>>(
-                styleTransferPtr, &runtime, jsCallInvoker);
-
-        return jsi::Object::createFromHostObject(runtime,
-                                                 styleTransferHostObject);
-      });
-}
-
 void RnExecutorchInstaller::injectJSIBindings(
-    jsi::Runtime *jsiRuntime,
-    const std::shared_ptr<react::CallInvoker> &jsCallInvoker,
+    jsi::Runtime *jsiRuntime, std::shared_ptr<react::CallInvoker> jsCallInvoker,
     FetchUrlFunc_t fetchDataFromUrl) {
   fetchUrlFunc = fetchDataFromUrl;
 
   jsiRuntime->global().setProperty(
       *jsiRuntime, "loadStyleTransfer",
-      loadStyleTransfer(jsiRuntime, jsCallInvoker));
+      RnExecutorchInstaller::loadModel<StyleTransfer>(jsiRuntime, jsCallInvoker,
+                                                      "loadStyleTransfer"));
+
+  jsiRuntime->global().setProperty(
+      *jsiRuntime, "loadImageSegmentation",
+      RnExecutorchInstaller::loadModel<ImageSegmentation>(
+          jsiRuntime, jsCallInvoker, "loadImageSegmentation"));
 }
 } // namespace rnexecutorch

android/src/main/java/com/swmansion/rnexecutorch/models/imageSegmentation/Constants.kt

@@ -7,23 +7,47 @@
 #include <ReactCommon/CallInvoker.h>
 #include <jsi/jsi.h>
 
+#include <rnexecutorch/host_objects/JsiConversions.h>
+#include <rnexecutorch/host_objects/ModelHostObject.h>
+
 namespace rnexecutorch {
 
 using FetchUrlFunc_t = std::function<std::vector<std::byte>(std::string)>;
+extern FetchUrlFunc_t fetchUrlFunc;
 
 using namespace facebook;
 
 class RnExecutorchInstaller {
 public:
   static void
   injectJSIBindings(jsi::Runtime *jsiRuntime,
-                    const std::shared_ptr<react::CallInvoker> &jsCallInvoker,
+                    std::shared_ptr<react::CallInvoker> jsCallInvoker,
                     FetchUrlFunc_t fetchDataFromUrl);
 
 private:
+  template <typename ModelT>
   static jsi::Function
-  loadStyleTransfer(jsi::Runtime *jsiRuntime,
-                    const std::shared_ptr<react::CallInvoker> &jsCallInvoker);
+  loadModel(jsi::Runtime *jsiRuntime,
+            std::shared_ptr<react::CallInvoker> jsCallInvoker,
+            const std::string &loadFunctionName) {
+    return jsi::Function::createFromHostFunction(
+        *jsiRuntime,
+        jsi::PropNameID::forAscii(*jsiRuntime, loadFunctionName.c_str()), 0,
+        [jsCallInvoker](jsi::Runtime &runtime, const jsi::Value &thisValue,
+                        const jsi::Value *args, size_t count) -> jsi::Value {
+          // We expect a single input -- the path to the model binary
+          assert(count == 1);
+          auto source = jsiconversion::getValue<std::string>(args[0], runtime);
+
+          auto modelImplementationPtr =
+              std::make_shared<ModelT>(source, &runtime);
+          auto modelHostObject = std::make_shared<ModelHostObject<ModelT>>(
+              modelImplementationPtr, jsCallInvoker);
+
+          return jsi::Object::createFromHostObject(runtime, modelHostObject);
+          return jsi::Value();
+        });
+  }
 };
 
 } // namespace rnexecutorch

android/src/main/java/com/swmansion/rnexecutorch/models/imageSegmentation/ImageSegmentationModel.kt

@@ -0,0 +1,18 @@
+#include "Numerical.h"
+
+#include <algorithm>
+#include <numeric>
+
+namespace rnexecutorch::numerical {
+void softmax(std::vector<float> &v) {
+  float max = *std::max_element(v.begin(), v.end());
+
+  for (float &x : v) {
+    x = std::exp(x - max);
+  }
+  float sum = std::accumulate(v.begin(), v.end(), 0.f);
+  for (float &x : v) {
+    x /= sum;
+  }
+}
+} // namespace rnexecutorch::numerical

common/rnexecutorch/RnExecutorchInstaller.cpp

@@ -0,0 +1,7 @@
+#pragma once
+
+#include <vector>
+
+namespace rnexecutorch::numerical {
+void softmax(std::vector<float> &v);
+}