main.rs

#![allow(unused_braces)]
use image::ImageReader;
use image::{DynamicImage, RgbImage};
use ndarray::{Array, Dim};
use std::fs;
use std::io::BufRead;

const IMG_PATH: &str = "fixture/images/dog.jpg";

wit_bindgen::generate!({
    path: "../../wit",
    world: "ml",
});

use self::wasi::nn::{
    graph::{Graph, GraphBuilder, load, ExecutionTarget, GraphEncoding},
    tensor::{Tensor, TensorData, TensorDimensions, TensorType},
};

/// Determine execution target from command-line argument
/// Usage: wasm_module [cpu|gpu|cuda]
fn get_execution_target() -> ExecutionTarget {
    let args: Vec<String> = std::env::args().collect();

    // First argument (index 0) is the program name, second (index 1) is the target
    // Ignore any arguments after index 1
    if args.len() >= 2 {
        match args[1].to_lowercase().as_str() {
            "gpu" | "cuda" => {
                println!("Using GPU (CUDA) execution target from argument");
                return ExecutionTarget::Gpu;
            }
            "cpu" => {
                println!("Using CPU execution target from argument");
                return ExecutionTarget::Cpu;
            }
            _ => {
                println!("Unknown execution target '{}', defaulting to CPU", args[1]);
            }
        }
    } else {
        println!("No execution target specified, defaulting to CPU");
        println!("Usage: <program> [cpu|gpu|cuda]");
    }

    ExecutionTarget::Cpu
}

fn main() {
    // Load the ONNX model - SqueezeNet 1.1-7
    // Full details: https://github.com/onnx/models/tree/main/vision/classification/squeezenet
    let model: GraphBuilder = fs::read("fixture/models/squeezenet1.1-7.onnx").unwrap();
    println!("Read ONNX model, size in bytes: {}", model.len());

    // Determine execution target
    let execution_target = get_execution_target();

    let graph = load(&[model], GraphEncoding::Onnx, execution_target).unwrap();
    println!("Loaded graph into wasi-nn with {:?} target", execution_target);

    let exec_context = Graph::init_execution_context(&graph).unwrap();
    println!("Created wasi-nn execution context.");

    // Load SquezeNet 1000 labels used for classification
    let labels = fs::read("fixture/labels/squeezenet1.1-7.txt").unwrap();
    let class_labels: Vec<String> = labels.lines().map(|line| line.unwrap()).collect();
    println!("Read ONNX Labels, # of labels: {}", class_labels.len());

    // Prepare WASI-NN tensor - Tensor data is always a bytes vector
    let dimensions: TensorDimensions = vec![1, 3, 224, 224];
    let data: TensorData = image_to_tensor(IMG_PATH.to_string(), 224, 224);
    let tensor = Tensor::new(
        &dimensions,
        TensorType::Fp32,
        &data,
    );
    let input_tensor = vec!(("data".to_string(), tensor));
    // Execute the inferencing
    let output_tensor_vec = exec_context.compute(input_tensor).unwrap();
    println!("Executed graph inference");

    let output_tensor = output_tensor_vec.iter().find_map(|(tensor_name, tensor)| {
        if tensor_name == "squeezenet0_flatten0_reshape0" {
            Some(tensor)
        } else {
            None
        }
    });
    let output_data = output_tensor.expect("No output tensor").data();

    println!("Retrieved output data with length: {}", output_data.len());
    let output_f32 = bytes_to_f32_vec(output_data);

    let output_shape = [1, 1000, 1, 1];
    let output_tensor = Array::from_shape_vec(output_shape, output_f32).unwrap();

    // Post-Processing requirement: compute softmax to inferencing output
    let exp_output = output_tensor.mapv(|x| x.exp());
    let sum_exp_output = exp_output.sum_axis(ndarray::Axis(1));
    let softmax_output = exp_output / &sum_exp_output;

    let mut sorted = softmax_output
        .axis_iter(ndarray::Axis(1))
        .enumerate()
        .into_iter()
        .map(|(i, v)| (i, v[Dim([0, 0, 0])]))
        .collect::<Vec<(_, _)>>();
    sorted.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap());

    for (index, probability) in sorted.iter().take(3) {
        println!(
            "Index: {} - Probability: {}",
            class_labels[*index], probability
        );
    }
}

pub fn bytes_to_f32_vec(data: Vec<u8>) -> Vec<f32> {
    let chunks: Vec<&[u8]> = data.chunks(4).collect();
    let v: Vec<f32> = chunks
        .into_iter()
        .map(|c| f32::from_le_bytes(c.try_into().unwrap()))
        .collect();

    v.into_iter().collect()
}

// Take the image located at 'path', open it, resize it to height x width, and then converts
// the pixel precision to FP32. The resulting BGR pixel vector is then returned.
fn image_to_tensor(path: String, height: u32, width: u32) -> Vec<u8> {
    let pixels = ImageReader::open(path).unwrap().decode().unwrap();
    let dyn_img: DynamicImage = pixels.resize_exact(width, height, image::imageops::Triangle);
    let bgr_img: RgbImage = dyn_img.to_rgb8();

    // Get an array of the pixel values
    let raw_u8_arr: &[u8] = &bgr_img.as_raw()[..];

    // Create an array to hold the f32 value of those pixels
    let bytes_required = raw_u8_arr.len() * 4;
    let mut u8_f32_arr: Vec<u8> = vec![0; bytes_required];

    // Normalizing values for the model
    let mean = [0.485, 0.456, 0.406];
    let std = [0.229, 0.224, 0.225];

    // Read the number as a f32 and break it into u8 bytes
    for i in 0..raw_u8_arr.len() {
        let u8_f32: f32 = raw_u8_arr[i] as f32;
        let rgb_iter = i % 3;

        // Normalize the pixel
        let norm_u8_f32: f32 = (u8_f32 / 255.0 - mean[rgb_iter]) / std[rgb_iter];

        // Convert it to u8 bytes and write it with new shape
        let u8_bytes = norm_u8_f32.to_ne_bytes();
        for j in 0..4 {
            u8_f32_arr[(raw_u8_arr.len() * 4 * rgb_iter / 3) + (i / 3) * 4 + j] = u8_bytes[j];
        }
    }

    return u8_f32_arr;
}