Merge pull request #94 from L-jasmine/feat/wasi_nn

feat: support wasi_nn

.devcontainer/devcontainer.json

@@ -0,0 +1,25 @@
+// For format details, see https://aka.ms/devcontainer.json. For config options, see the
+// README at: https://github.com/devcontainers/templates/tree/main/src/ubuntu
+{
+	"name": "Ubuntu",
+	// Or use a Dockerfile or Docker Compose file. More info: https://containers.dev/guide/dockerfile
+	"image": "mcr.microsoft.com/devcontainers/base:focal",
+	"features": {
+		"ghcr.io/devcontainers/features/rust:1": {}
+	}
+
+	// Features to add to the dev container. More info: https://containers.dev/features.
+	// "features": {},
+
+	// Use 'forwardPorts' to make a list of ports inside the container available locally.
+	// "forwardPorts": [],
+
+	// Use 'postCreateCommand' to run commands after the container is created.
+	// "postCreateCommand": "uname -a",
+
+	// Configure tool-specific properties.
+	// "customizations": {},
+
+	// Uncomment to connect as root instead. More info: https://aka.ms/dev-containers-non-root.
+	// "remoteUser": "root"
+}

.github/workflows/examples.yml

@@ -34,7 +34,11 @@ jobs:
       run: |
         VERSION=0.11.2
         curl -sSf https://raw.githubusercontent.com/WasmEdge/WasmEdge/master/utils/install.sh | sudo bash -s -- -e all --version=$VERSION --tf-version=$VERSION --tf-deps-version=$VERSION --tf-tools-version=$VERSION --image-version=$VERSION  -p /usr/local
-        
+        curl -sLO https://github.com/WasmEdge/WasmEdge/releases/download/0.12.0-alpha.1/WasmEdge-plugin-wasi_nn-tensorflowlite-0.12.0-alpha.1-ubuntu20.04_x86_64.tar.gz
+        tar -zxf WasmEdge-plugin-wasi_nn-tensorflowlite-*-ubuntu20.04_x86_64.tar.gz
+        rm -f WasmEdge-plugin-wasi_nn-tensorflowlite-*-ubuntu20.04_x86_64.tar.gz
+        sudo mv libwasmedgePluginWasiNN.so /usr/local/lib/wasmedge/
+
     - uses: actions/setup-node@v2
       with:
         node-version: '14'
@@ -153,6 +157,11 @@ jobs:
         cargo build --target wasm32-wasi --release --features=tensorflow
         wasmedge-tensorflow-lite --dir .:. target/wasm32-wasi/release/wasmedge_quickjs.wasm example_js/tensorflow_lite_demo/main.js
 
+    - name: WASI-NN example (TensorflowLite)
+      run: |
+        cargo build --target wasm32-wasi --release --features=wasi_nn
+        wasmedge --dir .:. target/wasm32-wasi/release/wasmedge_quickjs.wasm example_js/tensorflow_lite_demo/wasi_nn_main.js
+
     - name: Embed JS in Rust
       run: |
         cd examples/embed_js

Cargo.toml

@@ -9,15 +9,15 @@ readme = "README.md"
 documentation = "https://www.secondstate.io/articles/run-javascript-in-webassembly-with-wasmedge/"
 homepage = "https://www.secondstate.io/"
 edition = "2018"
-exclude = [
-    "example_js/*",
-    "examples/*",
-]
+exclude = ["example_js/*", "examples/*"]
 
 [dependencies]
 argparse = "0.2.2"
-image = { version = "0.23.6", default-features = false, features = ["jpeg", "png"], optional = true }
-imageproc = { version =  "0.22.0", optional = true }
+image = { version = "0.23.6", default-features = false, features = [
+    "jpeg",
+    "png",
+], optional = true }
+imageproc = { version = "0.22.0", optional = true }
 libc = "0.2"
 url = "2.2.2"
 lazy_static = "1.4"
@@ -27,4 +27,5 @@ encoding = "0.2"
 default = []
 img = ["image", "imageproc"]
 tensorflow = ["img"]
+wasi_nn = ["img"]
 cjs = []

example_js/tensorflow_lite_demo/wasi_nn_main.js

@@ -0,0 +1,36 @@
+import { Image } from 'image';
+import * as fs from 'fs';
+import { NnGraph, NnContext, TENSOR_TYPE_U8 } from 'wasi_nn';
+
+let img = new Image(__dirname + '/food.jpg');
+let img_rgb = img.to_rgb().resize(192, 192);
+let rgb_pix = img_rgb.pixels();
+
+let data = fs.readFileSync(__dirname + '/lite-model_aiy_vision_classifier_food_V1_1.tflite')
+let graph = new NnGraph([data.buffer], "tensorflowlite", "cpu");
+let context = new NnContext(graph);
+context.setInput(0, rgb_pix, [1, 192, 192, 3], TENSOR_TYPE_U8);
+context.compute();
+
+let output_view = new Uint8Array(2024);
+context.getOutput(0, output_view.buffer)
+
+let max = 0;
+let max_idx = 0;
+for (var i in output_view) {
+    let v = output_view[i];
+    if (v > max) {
+        max = v;
+        max_idx = i;
+    }
+}
+
+let label_file = fs.readFileSync(__dirname + '/aiy_food_V1_labelmap.txt', 'utf-8');
+let lables = label_file.split(/\r?\n/);
+
+let label = lables[max_idx]
+
+print('label:');
+print(label);
+print('confidence:');
+print(max / 255);

src/internal_module/mod.rs

@@ -8,3 +8,5 @@ pub mod os;
 #[cfg(feature = "tensorflow")]
 pub mod tensorflow_module;
 pub mod wasi_net_module;
+#[cfg(feature = "wasi_nn")]
+pub mod wasi_nn;

src/internal_module/tensorflow_module.rs

@@ -2,7 +2,7 @@ use crate::*;
 
 /// reference https://github.com/second-state/wasmedge_tensorflow_interface
 mod wasmedge_tensorflow {
-    /// wasmedge_tensorflow host functions.
+    // wasmedge_tensorflow host functions.
     #[link(wasm_import_module = "wasmedge_tensorflow")]
     extern "C" {
         pub fn wasmedge_tensorflow_create_session(model_buf: *const u8, model_buf_len: u32) -> u64;
@@ -37,7 +37,7 @@ mod wasmedge_tensorflow {
         pub fn wasmedge_tensorflow_clear_output(context: u64);
     }
 
-    /// wasmedge_tensorflowlite host functions.
+    // wasmedge_tensorflowlite host functions.
     #[link(wasm_import_module = "wasmedge_tensorflowlite")]
     extern "C" {
         pub fn wasmedge_tensorflowlite_create_session(
@@ -62,7 +62,7 @@ mod wasmedge_tensorflow {
         );
     }
 
-    /// wasmedge_image host helper functions.
+    // wasmedge_image host helper functions.
     #[link(wasm_import_module = "wasmedge_image")]
     extern "C" {
         pub fn wasmedge_image_load_jpg_to_rgb8(