fix: Validate input buffer type (#27)

* fix: Also validate input buffer type * fix: Use efficientdet
mrousavy · Jan 24, 2024 · e099d23 · e099d23
1 parent 99ccb41
commit e099d23
Show file tree

Hide file tree

Showing 6 changed files with 55 additions and 9 deletions.
diff --git a/cpp/TensorHelpers.cpp b/cpp/TensorHelpers.cpp
@@ -58,10 +58,34 @@ std::string dataTypeToString(TfLiteType dataType) {
     case kTfLiteInt4:
       return "int4";
     default:
+      [[unlikely]];
       return "invalid";
   }
 }
 
+TfLiteType getTFLDataTypeForTypedArrayKind(TypedArrayKind kind) {
+  switch (kind) {
+    case TypedArrayKind::Int8Array:
+      return kTfLiteInt8;
+    case TypedArrayKind::Int16Array:
+      return kTfLiteInt16;
+    case TypedArrayKind::Int32Array:
+      return kTfLiteInt32;
+    case TypedArrayKind::Uint8Array:
+      return kTfLiteUInt8;
+    case TypedArrayKind::Uint8ClampedArray:
+      return kTfLiteUInt8;
+    case TypedArrayKind::Uint16Array:
+      return kTfLiteUInt16;
+    case TypedArrayKind::Uint32Array:
+      return kTfLiteUInt32;
+    case TypedArrayKind::Float32Array:
+      return kTfLiteFloat32;
+    case TypedArrayKind::Float64Array:
+      return kTfLiteFloat64;
+  }
+}
+
 size_t TensorHelpers::getTFLTensorDataTypeSize(TfLiteType dataType) {
   switch (dataType) {
     case kTfLiteFloat32:
@@ -195,8 +219,20 @@ void TensorHelpers::updateJSBufferFromTensor(jsi::Runtime& runtime, TypedArrayBa
 
 void TensorHelpers::updateTensorFromJSBuffer(jsi::Runtime& runtime, TfLiteTensor* tensor,
                                              TypedArrayBase& jsBuffer) {
-  auto name = std::string(TfLiteTensorName(tensor));
-  auto buffer = jsBuffer.getBuffer(runtime);
+#if DEBUG
+  TypedArrayKind kind = jsBuffer.getKind(runtime);
+  TfLiteType receivedType = getTFLDataTypeForTypedArrayKind(kind);
+  TfLiteType expectedType = TfLiteTensorType(tensor);
+  if (receivedType != expectedType) {
+    [[unlikely]];
+    throw std::runtime_error("Invalid input type! Model expected " +
+                             dataTypeToString(expectedType) + ", but received " +
+                             dataTypeToString(receivedType) + "!");
+  }
+#endif
+
+  std::string name = TfLiteTensorName(tensor);
+  jsi::ArrayBuffer buffer = jsBuffer.getBuffer(runtime);
 
 #if DEBUG
   int inputBufferSize = buffer.size(runtime);

diff --git a/cpp/TensorflowPlugin.cpp b/cpp/TensorflowPlugin.cpp
@@ -202,7 +202,15 @@ TensorflowPlugin::getOutputArrayForTensor(jsi::Runtime& runtime, const TfLiteTen
 
 void TensorflowPlugin::copyInputBuffers(jsi::Runtime& runtime, jsi::Object inputValues) {
   // Input has to be array in input tensor size
-  auto array = inputValues.asArray(runtime);
+#if DEBUG
+  if (!inputValues.isArray(runtime)) {
+    [[unlikely]];
+    throw std::runtime_error(
+        "TFLite: Input Values must be an array, one item for each input tensor!");
+  }
+#endif
+
+  jsi::Array array = inputValues.asArray(runtime);
   size_t count = array.size(runtime);
   if (count != TfLiteInterpreterGetInputTensorCount(_interpreter)) {
     [[unlikely]];
@@ -212,8 +220,8 @@ void TensorflowPlugin::copyInputBuffers(jsi::Runtime& runtime, jsi::Object input
 
   for (size_t i = 0; i < count; i++) {
     TfLiteTensor* tensor = TfLiteInterpreterGetInputTensor(_interpreter, i);
-    auto value = array.getValueAtIndex(runtime, i);
-    auto inputBuffer = getTypedArray(runtime, value.asObject(runtime));
+    jsi::Value value = array.getValueAtIndex(runtime, i);
+    TypedArrayBase inputBuffer = getTypedArray(runtime, value.asObject(runtime));
     TensorHelpers::updateTensorFromJSBuffer(runtime, tensor, inputBuffer);
   }
 }

diff --git a/example/assets/efficientdet.tflite b/example/assets/efficientdet.tflite
diff --git a/example/assets/object_detector.tflite b/example/assets/object_detector.tflite
diff --git a/example/assets/smartreply_1_default_1.tflite b/example/assets/smartreply_1_default_1.tflite
diff --git a/example/src/App.tsx b/example/src/App.tsx
@@ -30,7 +30,8 @@ export default function App(): React.ReactNode {
   const { hasPermission, requestPermission } = useCameraPermission()
   const device = useCameraDevice('back')
 
-  const model = useTensorflowModel(require('../assets/object_detector.tflite'))
+  // from https://www.kaggle.com/models/tensorflow/efficientdet/frameworks/tfLite
+  const model = useTensorflowModel(require('../assets/efficientdet.tflite'))
   const actualModel = model.state === 'loaded' ? model.model : undefined
 
   React.useEffect(() => {
@@ -51,14 +52,15 @@ export default function App(): React.ReactNode {
       console.log(`Running inference on ${frame}`)
       const resized = resize(frame, {
         size: {
-          width: 640,
-          height: 640,
+          width: 320,
+          height: 320,
         },
         pixelFormat: 'rgb-uint8',
       })
       const typedArray = new Uint8Array(resized)
       const result = actualModel.runSync([typedArray])
-      console.log('Result: ' + result.length)
+      const num_detections = result[3]?.[0] ?? 0
+      console.log('Result: ' + num_detections)
     },
     [actualModel]
   )