-
Notifications
You must be signed in to change notification settings - Fork 23
/
person_detection.rs
84 lines (66 loc) · 2.37 KB
/
person_detection.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
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
//! person_detection example
//!
use tfmicro::{MicroInterpreter, Model, MutableOpResolver};
use itertools::Itertools;
use log::info;
#[test]
fn person_detection() {
env_logger::init();
info!("---- Starting tensorflow micro example: person_detection");
// Include trained model and test datasets
let model =
include_bytes!("../examples/models/person_detection_grayscale.tflite");
let no_person = include_bytes!(
"../examples/models/no_person_image_data_grayscale.data"
);
let person =
include_bytes!("../examples/models/person_image_data_grayscale.data");
// Map the model into a usable data structure. This doesn't involve
// any copying or parsing, it's a very lightweight operation.
let model = Model::from_buffer(&model[..]).unwrap();
// Create memory area for input, output and intermediate arrays
const TENSOR_ARENA_SIZE: usize = 93 * 1024;
let mut tensor_arena: [u8; TENSOR_ARENA_SIZE] = [0; TENSOR_ARENA_SIZE];
let micro_op_resolver = MutableOpResolver::empty()
.depthwise_conv_2d()
.conv_2d()
.average_pool_2d();
// Build an interpreter to run the model with
let mut interpreter =
MicroInterpreter::new(&model, micro_op_resolver, &mut tensor_arena[..])
.unwrap();
// Check properties of the input sensor
interpreter.input(0, person).unwrap();
assert_eq!([1, 96, 96, 1], interpreter.input_info(0).dims);
info!("Created setup");
// -------- 'person' example ------------
interpreter.invoke().unwrap();
// get output for 'person'
let output_tensor = interpreter.output(0);
assert_eq!(
[1, 1, 1, 3],
output_tensor.info().dims,
"Dimensions of output tensor"
);
assert_eq!(
1,
output_tensor.as_data::<u8>().iter().position_max().unwrap()
);
info!("---- Person output correct!");
// ------- 'no person' example ----------
interpreter.input(0, no_person).unwrap();
interpreter.invoke().unwrap();
// get output for 'no person'
let output_tensor = interpreter.output(0);
assert_eq!(
[1, 1, 1, 3],
output_tensor.info().dims,
"Dimensions of output tensor"
);
assert_eq!(
2,
output_tensor.as_data::<u8>().iter().position_max().unwrap()
);
info!("---- No-person output correct!");
info!("---- Done");
}