Skip to content

Commit 1942a29

Browse files
gusye1234ye jianbai
authored andcommitted
feat: add tf builder
Signed-off-by: Jianbai Ye <jianbaiye@outlook.com>
1 parent 62d8787 commit 1942a29

File tree

3 files changed

+7
-5
lines changed

3 files changed

+7
-5
lines changed

tf-mobilenet-image/rust/tf-mobilenet/Cargo.toml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@ publish = false
88

99
[dependencies]
1010
image = { version = "0.23.14", default-features = false, features = ["gif", "jpeg", "ico", "png", "pnm", "tga", "tiff", "webp", "bmp", "hdr", "dxt", "dds", "farbfeld"] }
11-
wasi-nn = { version = "0.2.1", path="../../../../wasmedge-wasi-nn/rust"}
11+
wasi-nn = { version = "0.2.1"}
1212

1313
# This crate is built with the wasm32-wasi target, so it's separate
1414
# from the main Wasmtime build, so use this directive to exclude it

tf-mobilenet-image/rust/tf-mobilenet/src/main.rs

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -21,11 +21,13 @@ fn infer_image() {
2121
let image_name: &str = &args[2];
2222

2323
let weights = fs::read(model_bin_name).unwrap();
24+
let tagset = "";
25+
let signature = "image_classifier";
2426
println!("Read graph weights, size in bytes: {}", weights.len());
2527

2628
let graph = unsafe {
2729
wasi_nn::load(
28-
&[&weights],
30+
&[&weights, &tagset.as_bytes(), &signature.as_bytes()],
2931
wasi_nn::GRAPH_ENCODING_TENSORFLOW, // encoding for tflite
3032
wasi_nn::EXECUTION_TARGET_CPU,
3133
)
@@ -45,7 +47,7 @@ fn infer_image() {
4547
data: &tensor_data,
4648
};
4749
unsafe {
48-
wasi_nn::set_input_by_string(context, "hub_input/images", tensor).unwrap();
50+
wasi_nn::set_input(context, 0, tensor).unwrap();
4951
}
5052
println!("Set input done");
5153
// Execute the inference.
@@ -56,9 +58,9 @@ fn infer_image() {
5658
// Retrieve the output.
5759
let mut output_buffer = vec![0f32; 965];
5860
unsafe {
59-
wasi_nn::get_output_by_string(
61+
wasi_nn::get_output(
6062
context,
61-
"prediction",
63+
0,
6264
&mut output_buffer[..] as *mut [f32] as *mut u8,
6365
(output_buffer.len() * 4).try_into().unwrap(),
6466
)
Binary file not shown.

0 commit comments

Comments
 (0)