-
Notifications
You must be signed in to change notification settings - Fork 102
/
lib.rs
64 lines (57 loc) · 2.25 KB
/
lib.rs
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
use wasm_bindgen::prelude::*;
use wasmedge_tensorflow_interface;
use image::{GenericImageView, Pixel};
use imageproc::drawing::draw_hollow_rect_mut;
use imageproc::rect::Rect;
use std::str;
use std::time::{Instant};
#[wasm_bindgen]
pub fn infer(image_data: &[u8]) -> Vec<u8> {
let start = Instant::now();
let mut img = image::load_from_memory(image_data).unwrap();
let mut flat_img: Vec<f32> = Vec::new();
for (_x, _y, rgb) in img.pixels() {
flat_img.push(rgb[2] as f32);
flat_img.push(rgb[1] as f32);
flat_img.push(rgb[0] as f32);
}
println!("Loaded image in ... {:?}", start.elapsed());
let model_data: &[u8] = include_bytes!("mtcnn.pb");
let mut session = wasmedge_tensorflow_interface::Session::new(model_data, wasmedge_tensorflow_interface::ModelType::TensorFlow);
session.add_input("min_size", &[20.0f32], &[])
.add_input("thresholds", &[0.6f32, 0.7f32, 0.7f32], &[3])
.add_input("factor", &[0.709f32], &[])
.add_input("input", &flat_img, &[img.height().into(), img.width().into(), 3])
.add_output("box")
.add_output("prob")
.run();
let res_vec: Vec<f32> = session.get_output("box");
// Parse results.
let mut iter = 0;
let mut box_vec: Vec<[f32; 4]> = Vec::new();
while (iter * 4) < res_vec.len() {
box_vec.push([
res_vec[4 * iter + 1], // x1
res_vec[4 * iter], // y1
res_vec[4 * iter + 3], // x2
res_vec[4 * iter + 2], // y2
]);
iter += 1;
}
println!("Parsed results in ... {:?}", start.elapsed());
println!("Drawing box: {} results ...", box_vec.len());
let line = Pixel::from_slice(&[0, 255, 0, 0]);
for i in 0..box_vec.len() {
let xy = box_vec[i];
let x1: i32 = xy[0] as i32;
let y1: i32 = xy[1] as i32;
let x2: i32 = xy[2] as i32;
let y2: i32 = xy[3] as i32;
let rect = Rect::at(x1, y1).of_size((x2 - x1) as u32, (y2 - y1) as u32);
draw_hollow_rect_mut(&mut img, rect, *line);
}
let mut buf = Vec::new();
img.write_to(&mut buf, image::ImageOutputFormat::Jpeg(80u8)).expect("Unable to write");
println!("Drawn on image in ... {:?}", start.elapsed());
return buf;
}