fix: abstract loading images to tensors

JakubGonera · JakubGonera · commit ae5a90f583f9 · 2025-06-06T13:45:22.000+02:00
diff --git a/packages/react-native-executorch/common/rnexecutorch/data_processing/ImageProcessing.cpp b/packages/react-native-executorch/common/rnexecutorch/data_processing/ImageProcessing.cpp
@@ -111,20 +111,42 @@ cv::Mat readImage(const std::string &imageURI) {
     throw std::runtime_error("Read image error: invalid argument");
   }
 
-  cv::cvtColor(image, image, cv::COLOR_BGR2RGB);
   return image;
 }
 
-TensorPtr getTensorFromMatrix(const std::vector<int32_t> &sizes,
+TensorPtr getTensorFromMatrix(const std::vector<int32_t> &tensorDim,
                               const cv::Mat &matrix) {
   std::vector<float> inputVector = colorMatToVector(matrix);
-  return executorch::extension::make_tensor_ptr(sizes, inputVector);
+  return executorch::extension::make_tensor_ptr(tensorDim, inputVector);
 }
 
 cv::Mat getMatrixFromTensor(cv::Size size, const Tensor &tensor) {
   auto resultData = static_cast<const float *>(tensor.const_data_ptr());
   return bufferToColorMat(std::span<const float>(resultData, tensor.numel()),
                           size);
 }
+
+std::pair<TensorPtr, cv::Size>
+readImageToTensor(const std::string &path,
+                  const std::vector<int32_t> &tensorDim) {
+  cv::Mat input = imageprocessing::readImage(path);
+  cv::cvtColor(input, input, cv::COLOR_BGR2RGB);
+  cv::Size imageSize = input.size();
+
+  if (tensorDim.size() < 2) {
+    char errorMessage[100];
+    std::snprintf(errorMessage, sizeof(errorMessage),
+                  "Unexpected model input size, expected at least 2 dimentions "
+                  "but got: %zu.",
+                  tensorDim.size());
+    throw std::runtime_error(errorMessage);
+  }
+  cv::Size tensorSize = cv::Size(tensorDim[tensorDim.size() - 1],
+                                 tensorDim[tensorDim.size() - 2]);
+
+  cv::resize(input, input, tensorSize);
+
+  return {imageprocessing::getTensorFromMatrix(tensorDim, input), imageSize};
+}
 } // namespace imageprocessing
 } // namespace rnexecutorch
diff --git a/packages/react-native-executorch/common/rnexecutorch/data_processing/ImageProcessing.h b/packages/react-native-executorch/common/rnexecutorch/data_processing/ImageProcessing.h
@@ -1,12 +1,15 @@
 #pragma once
 
-#include <executorch/extension/tensor/tensor.h>
-#include <executorch/extension/tensor/tensor_ptr.h>
-#include <opencv2/opencv.hpp>
+#include <optional>
 #include <span>
 #include <string>
 #include <vector>
 
+#include <executorch/extension/tensor/tensor.h>
+#include <executorch/extension/tensor/tensor_ptr.h>
+
+#include <opencv2/opencv.hpp>
+
 namespace rnexecutorch::imageprocessing {
 using executorch::aten::Tensor;
 using executorch::extension::TensorPtr;
@@ -22,8 +25,13 @@ cv::Mat bufferToColorMat(const std::span<const float> &buffer,
                          cv::Size matSize);
 std::string saveToTempFile(const cv::Mat &image);
 cv::Mat readImage(const std::string &imageURI);
-TensorPtr getTensorFromMatrix(const std::vector<int32_t> &sizes,
+TensorPtr getTensorFromMatrix(const std::vector<int32_t> &tensorDim,
                               const cv::Mat &mat);
 cv::Mat getMatrixFromTensor(cv::Size size, const Tensor &tensor);
+/// @brief Read image, resize it and copy it to an ET tensor to store it.
+/// @return Returns a tensor pointer and the original size of the image.
+std::pair<TensorPtr, cv::Size>
+readImageToTensor(const std::string &path,
+                  const std::vector<int32_t> &tensorDim);
 
 } // namespace rnexecutorch::imageprocessing
diff --git a/packages/react-native-executorch/common/rnexecutorch/models/classification/Classification.cpp b/packages/react-native-executorch/common/rnexecutorch/models/classification/Classification.cpp
@@ -30,7 +30,8 @@ Classification::Classification(const std::string &modelSource,
 
 std::unordered_map<std::string_view, float>
 Classification::forward(std::string imageSource) {
-  auto tensor = preprocess(imageSource);
+  auto tensor =
+      imageprocessing::readImageToTensor(imageSource, getInputShape()[0]).first;
 
   auto forwardResult = forwardET(tensor);
   if (!forwardResult.ok()) {
@@ -42,13 +43,6 @@ Classification::forward(std::string imageSource) {
   return postprocess(forwardResult->at(0).toTensor());
 }
 
-TensorPtr Classification::preprocess(const std::string &imageSource) {
-  cv::Mat image = imageprocessing::readImage(imageSource);
-  cv::resize(image, image, modelImageSize);
-
-  return imageprocessing::getTensorFromMatrix(getInputShape()[0], image);
-}
-
 std::unordered_map<std::string_view, float>
 Classification::postprocess(const Tensor &tensor) {
   std::span<const float> resultData(
diff --git a/packages/react-native-executorch/common/rnexecutorch/models/classification/Classification.h b/packages/react-native-executorch/common/rnexecutorch/models/classification/Classification.h
@@ -18,7 +18,6 @@ class Classification : public BaseModel {
   std::unordered_map<std::string_view, float> forward(std::string imageSource);
 
 private:
-  TensorPtr preprocess(const std::string &imageSource);
   std::unordered_map<std::string_view, float> postprocess(const Tensor &tensor);
 
   cv::Size modelImageSize{0, 0};
diff --git a/packages/react-native-executorch/common/rnexecutorch/models/image_segmentation/ImageSegmentation.cpp b/packages/react-native-executorch/common/rnexecutorch/models/image_segmentation/ImageSegmentation.cpp
@@ -37,7 +37,8 @@ std::shared_ptr<jsi::Object>
 ImageSegmentation::forward(std::string imageSource,
                            std::set<std::string, std::less<>> classesOfInterest,
                            bool resize) {
-  auto [inputTensor, originalSize] = preprocess(imageSource);
+  auto [inputTensor, originalSize] =
+      imageprocessing::readImageToTensor(imageSource, getInputShape()[0]);
 
   auto forwardResult = forwardET(inputTensor);
   if (!forwardResult.ok()) {
@@ -50,19 +51,6 @@ ImageSegmentation::forward(std::string imageSource,
                      classesOfInterest, resize);
 }
 
-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::shared_ptr<jsi::Object> ImageSegmentation::postprocess(
     const Tensor &tensor, cv::Size originalSize,
     std::set<std::string, std::less<>> classesOfInterest, bool resize) {
diff --git a/packages/react-native-executorch/common/rnexecutorch/models/image_segmentation/ImageSegmentation.h b/packages/react-native-executorch/common/rnexecutorch/models/image_segmentation/ImageSegmentation.h
@@ -26,7 +26,6 @@ class ImageSegmentation : public BaseModel {
           std::set<std::string, std::less<>> classesOfInterest, bool resize);
 
 private:
-  std::pair<TensorPtr, cv::Size> preprocess(const std::string &imageSource);
   std::shared_ptr<jsi::Object>
   postprocess(const Tensor &tensor, cv::Size originalSize,
               std::set<std::string, std::less<>> classesOfInterest,
diff --git a/packages/react-native-executorch/common/rnexecutorch/models/object_detection/ObjectDetection.cpp b/packages/react-native-executorch/common/rnexecutorch/models/object_detection/ObjectDetection.cpp
@@ -25,16 +25,6 @@ ObjectDetection::ObjectDetection(
                             modelInputShape[modelInputShape.size() - 2]);
 }
 
-std::pair<TensorPtr, cv::Size>
-ObjectDetection::preprocess(const std::string &imageSource) {
-  cv::Mat image = imageprocessing::readImage(imageSource);
-  auto originalSize = image.size();
-  cv::resize(image, image, modelImageSize);
-
-  return {imageprocessing::getTensorFromMatrix(getInputShape()[0], image),
-          originalSize};
-}
-
 std::vector<Detection>
 ObjectDetection::postprocess(const std::vector<EValue> &tensors,
                              cv::Size originalSize, double detectionThreshold) {
@@ -75,7 +65,8 @@ ObjectDetection::postprocess(const std::vector<EValue> &tensors,
 
 std::vector<Detection> ObjectDetection::forward(std::string imageSource,
                                                 double detectionThreshold) {
-  auto [tensor, originalSize] = preprocess(imageSource);
+  auto [tensor, originalSize] =
+      imageprocessing::readImageToTensor(imageSource, getInputShape()[0]);
 
   auto forwardResult = forwardET(tensor);
   if (!forwardResult.ok()) {
diff --git a/packages/react-native-executorch/common/rnexecutorch/models/object_detection/ObjectDetection.h b/packages/react-native-executorch/common/rnexecutorch/models/object_detection/ObjectDetection.h
@@ -21,7 +21,6 @@ class ObjectDetection : public BaseModel {
                                  double detectionThreshold);
 
 private:
-  std::pair<TensorPtr, cv::Size> preprocess(const std::string &imageSource);
   std::vector<Detection> postprocess(const std::vector<EValue> &tensors,
                                      cv::Size originalSize,
                                      double detectionThreshold);
diff --git a/packages/react-native-executorch/common/rnexecutorch/models/style_transfer/StyleTransfer.cpp b/packages/react-native-executorch/common/rnexecutorch/models/style_transfer/StyleTransfer.cpp
@@ -33,16 +33,6 @@ StyleTransfer::StyleTransfer(const std::string &modelSource,
                             modelInputShape[modelInputShape.size() - 2]);
 }
 
-std::pair<TensorPtr, cv::Size>
-StyleTransfer::preprocess(const std::string &imageSource) {
-  cv::Mat image = imageprocessing::readImage(imageSource);
-  auto originalSize = image.size();
-  cv::resize(image, image, modelImageSize);
-
-  return {imageprocessing::getTensorFromMatrix(getInputShape()[0], image),
-          originalSize};
-}
-
 std::string StyleTransfer::postprocess(const Tensor &tensor,
                                        cv::Size originalSize) {
   cv::Mat mat = imageprocessing::getMatrixFromTensor(modelImageSize, tensor);
@@ -52,7 +42,8 @@ std::string StyleTransfer::postprocess(const Tensor &tensor,
 }
 
 std::string StyleTransfer::forward(std::string imageSource) {
-  auto [tensor, originalSize] = preprocess(imageSource);
+  auto [tensor, originalSize] =
+      imageprocessing::readImageToTensor(imageSource, getInputShape()[0]);
 
   auto forwardResult = forwardET(tensor);
   if (!forwardResult.ok()) {
diff --git a/packages/react-native-executorch/common/rnexecutorch/models/style_transfer/StyleTransfer.h b/packages/react-native-executorch/common/rnexecutorch/models/style_transfer/StyleTransfer.h
@@ -22,7 +22,6 @@ class StyleTransfer : public BaseModel {
   std::string forward(std::string imageSource);
 
 private:
-  std::pair<TensorPtr, cv::Size> preprocess(const std::string &imageSource);
   std::string postprocess(const Tensor &tensor, cv::Size originalSize);
 
   cv::Size modelImageSize{0, 0};
