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,6 +1,7 @@
 #include "RnExecutorchInstaller.h"
 
 #include <rnexecutorch/host_objects/JsiConversions.h>
+#include <rnexecutorch/models/image_segmentation/ImageSegmentation.h>
 #include <rnexecutorch/models/style_transfer/StyleTransfer.h>
 
 namespace rnexecutorch {
@@ -19,5 +20,10 @@ void RnExecutorchInstaller::injectJSIBindings(
       *jsiRuntime, "loadStyleTransfer",
       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

@@ -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

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

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

common/rnexecutorch/RnExecutorchInstaller.cpp

@@ -1,8 +1,10 @@
 #pragma once
 
-#include <jsi/jsi.h>
+#include <set>
 #include <type_traits>
 
+#include <jsi/jsi.h>
+
 namespace rnexecutorch::jsiconversion {
 
 using namespace facebook;
@@ -43,6 +45,25 @@ getValue<std::vector<std::string>>(const jsi::Value &val,
   return result;
 }
 
+// Set with heterogenerous look-up (adding std::less<> enables querying
+// with std::string_view)
+template <>
+inline std::set<std::string, std::less<>>
+getValue<std::set<std::string, std::less<>>>(const jsi::Value &val,
+                                             jsi::Runtime &runtime) {
+  // C++ set from JS array
+
+  jsi::Array array = val.asObject(runtime).asArray(runtime);
+  size_t length = array.size(runtime);
+  std::set<std::string, std::less<>> result;
+
+  for (size_t i = 0; i < length; ++i) {
+    jsi::Value element = array.getValueAtIndex(runtime, i);
+    result.insert(getValue<std::string>(element, runtime));
+  }
+  return result;
+}
+
 // Conversion from C++ types to jsi --------------------------------------------
 
 // Implementation functions might return any type, but in a promise we can only
@@ -54,10 +75,6 @@ inline jsi::Value getJsiValue(jsi::Value &&value, jsi::Runtime &runtime) {
   return std::move(value);
 }
 
-inline jsi::Value getJsiValue(jsi::Object &&value, jsi::Runtime &runtime) {
-  return jsi::Value(std::move(value));
-}
-
 inline jsi::Value getJsiValue(const std::string &str, jsi::Runtime &runtime) {
   return jsi::String::createFromAscii(runtime, str);
 }

common/rnexecutorch/data_processing/Numerical.cpp

@@ -0,0 +1,29 @@
+#pragma once
+
+#include <jsi/jsi.h>
+
+namespace rnexecutorch {
+
+using namespace facebook;
+
+class OwningArrayBuffer : public jsi::MutableBuffer {
+public:
+  OwningArrayBuffer(const size_t size) : size_(size) {
+    data_ = new uint8_t[size];
+  }
+  ~OwningArrayBuffer() override { delete[] data_; }
+
+  OwningArrayBuffer(const OwningArrayBuffer &) = delete;
+  OwningArrayBuffer(OwningArrayBuffer &&) = delete;
+  OwningArrayBuffer &operator=(const OwningArrayBuffer &) = delete;
+  OwningArrayBuffer &operator=(OwningArrayBuffer &&) = delete;
+
+  [[nodiscard]] size_t size() const override { return size_; }
+  uint8_t *data() override { return data_; }
+
+private:
+  uint8_t *data_;
+  const size_t size_;
+};
+
+} // namespace rnexecutorch

common/rnexecutorch/data_processing/Numerical.h

@@ -14,7 +14,10 @@ class BaseModel {
 
 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;
 };
-
 } // namespace rnexecutorch