1 use anyhow::Result;
2 use std::time::Instant;
3 use wasi_nn::{Graph, TensorType};
4 
5 /// Run a wasi-nn inference using a simple classifier model (single input,
6 /// single output).
7 pub fn classify(graph: Graph, tensor: Vec<u8>) -> Result<Vec<f32>> {
8     let mut context = graph.init_execution_context()?;
9     println!(
10         "[nn] created wasi-nn execution context with ID: {}",
11         context
12     );
13 
14     // Many classifiers have a single input; currently, this test suite also
15     // uses tensors of the same shape, though this is not usually the case.
16     context.set_input(0, TensorType::F32, &[1, 3, 224, 224], &tensor)?;
17     println!("[nn] set input tensor: {} bytes", tensor.len());
18 
19     let before = Instant::now();
20     context.compute()?;
21     println!(
22         "[nn] executed graph inference in {} ms",
23         before.elapsed().as_millis()
24     );
25 
26     // Many classifiers emit probabilities as floating point values; here we
27     // convert the raw bytes to `f32` knowing all models used here use that
28     // type.
29     let mut output_buffer = vec![0u8; 1001 * std::mem::size_of::<f32>()];
30     let num_bytes = context.get_output(0, &mut output_buffer)?;
31     println!("[nn] retrieved output tensor: {} bytes", num_bytes);
32     let output: Vec<f32> = output_buffer[..num_bytes]
33         .chunks(4)
34         .map(|b| f32::from_le_bytes([b[0], b[1], b[2], b[3]]))
35         .collect();
36     Ok(output)
37 }
38 
39 /// Sort some classification probabilities.
40 ///
41 /// Many classification models output a buffer of probabilities for each class,
42 /// placing the match probability for each class at the index for that class
43 /// (the probability of class `N` is stored at `probabilities[N]`).
44 pub fn sort_results(probabilities: &[f32]) -> Vec<InferenceResult> {
45     // It is unclear why the MobileNet output indices are "off by one" but the
46     // `.skip(1)` below seems necessary to get results that make sense (e.g. 763
47     // = "revolver" vs 762 = "restaurant").
48     let mut results: Vec<InferenceResult> = probabilities
49         .iter()
50         .skip(1)
51         .enumerate()
52         .map(|(c, p)| InferenceResult(c, *p))
53         .collect();
54     results.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap());
55     results
56 }
57 
58 // A wrapper for class ID and match probabilities.
59 #[derive(Debug, PartialEq)]
60 pub struct InferenceResult(usize, f32);
61 impl InferenceResult {
62     pub fn class_id(&self) -> usize {
63         self.0
64     }
65     pub fn probability(&self) -> f32 {
66         self.1
67     }
68 }
69