Running ML models in a game (and in Wasm!)
Tags: rust, bevy, webassembly, onnx
Published: 2021-02-19 21:15:31

Using [Tract](https://github.com/sonos/tract) and [Bevy](https://bevyengine.org), guess the number being drawn.

TL;DR: [Try it here](https://vleue.itch.io/bevmnist-poc) - [Read the code here](https://github.com/vleue/bevmnist)

One of my colleague recently starred [Tract github repository](https://github.com/sonos/tract) and got me wondering how easy it would be to use. I know how to create a [ONNX model](https://onnx.ai) with [PyTorch Lightning](https://www.pytorchlightning.ai), or even that there are [pretrained models](https://github.com/onnx/models) available. Spoilers: it's very easy to integrate!

The last couple of month, I've been doing game jams, [Ludum Dare 47](https://vleue.itch.io/kmanb) and [Game Off 2020](https://vleue.itch.io/wotm) with [Bevy](https://bevyengine.org), so I wanted to check If I could easily use Tract in a game made with Bevy and build for a Wasm target.

![guessing all digits](https://dev-to-uploads.s3.amazonaws.com/uploads/articles/oapj28leslekhx97hhhk.gif)

## Running a ONNX model with Tract

### Tract example

Using Tract to run a ONNX model is fearly easy, and the [example](https://github.com/sonos/tract/blob/main/examples/onnx-mobilenet-v2/src/main.rs) provided is great.

First, you load the model, specifying it's input:

```rust
let model = tract_onnx::onnx()
    // load the model
    .model_for_path("mobilenetv2-1.0.onnx")?
    // specify input type and shape
    .with_input_fact(0, InferenceFact::dt_shape(f32::datum_type(), tvec!(1, 3, 224, 224)))?
    // optimize the model
    .into_optimized()?
    // make the model runnable and fix its inputs and outputs
    .into_runnable()?;
```

Then take an image and transform it to an array with the expected shape and normalize the values:

```rust
// open image, resize it and make a Tensor out of it
let image = image::open("grace_hopper.jpg").unwrap().to_rgb8();
let resized =
    image::imageops::resize(&image, 224, 224, ::image::imageops::FilterType::Triangle);
let image: Tensor = tract_ndarray::Array4::from_shape_fn((1, 3, 224, 224), |(_, c, y, x)| {
    let mean = [0.485, 0.456, 0.406][c];
    let std = [0.229, 0.224, 0.225][c];
    (resized[(x as _, y as _)][c] as f32 / 255.0 - mean) / std
})
.into();
```

And finally run the model and get the result with the best score:

```rust
// run the model on the input
let result = model.run(tvec!(image))?;

// find and display the max value with its index
let best = result[0]
    .to_array_view::<f32>()?
    .iter()
    .cloned()
    .zip(2..)
    .max_by(|a, b| a.0.partial_cmp(&b.0).unwrap());
println!("result: {:?}", best);
```

### With MNIST

For MNIST, I use the model from [ONNX MNIST](https://github.com/onnx/models/tree/master/vision/classification/mnist). It takes an image of 28px by 28px as input, so its shape is `(1, 1, 28, 28)`.

Its output is an array of 10 float numbers, representing the score of each digit. I can get the digit with the best score:

```rust
let result = model.model.run(tvec!(image)).unwrap();

if let Some((value, score)) = result[0]
    .to_array_view::<f32>()
    .unwrap()
    .iter()
    .cloned()
    .enumerate()
    .max_by(|a, b| a.1.partial_cmp(&b.1).unwrap())
{
    if score > 10. {
        println!("{:?}", value);
    }
}
```


## Game setup

### Loading the model

To load the model from an `.onnx` file, I need to create a custom asset loader for this format. Following [bevy custom asset loader example](https://github.com/bevyengine/bevy/blob/master/examples/asset/custom_asset.rs), I first declare my [asset type `OnnxModel`](https://github.com/vleue/bevmnist/blob/57b51cd5121f88af04e1defd40e9b8e1cb9c9623/src/main.rs#L14)

```rust
#[derive(Debug, TypeUuid)]
#[uuid = "578fae90-a8de-41ab-a4dc-3aca66a31eed"]
pub struct OnnxModel {
    pub model: SimplePlan<
        TypedFact,
        Box<dyn TypedOp>,
        tract_onnx::prelude::Graph<TypedFact, Box<dyn TypedOp>>,
    >,
}
```

And then I implement `AssetLoader` for [`OnnxModelLoader`](https://github.com/vleue/bevmnist/blob/57b51cd5121f88af04e1defd40e9b8e1cb9c9623/src/main.rs#L25){:target=_"blank"}. Even though I know the input shape of my model, I do not call `with_input_fact` in the loader to be independent on the model loaded.

```rust
#[derive(Default)]
pub struct OnnxModelLoader;

impl AssetLoader for OnnxModelLoader {
    fn load<'a>(
        &'a self,
        mut bytes: &'a [u8],
        load_context: &'a mut LoadContext,
    ) -> BoxedFuture<'a, Result<(), anyhow::Error>> {
        Box::pin(async move {
            let model = tract_onnx::onnx()
                .model_for_read(&mut bytes)
                .unwrap()
                .into_optimized()?
                .into_runnable()?;

            load_context.set_default_asset(LoadedAsset::new(OnnxModel { model }));
            Ok(())
        })
    }

    fn extensions(&self) -> &[&str] {
        &["onnx"]
    }
}
```

I create a `struct` to hold the `Handle` to the model so that I can reuse the loaded model without needing to reload it every time

```rust
struct State {
    model: Handle<OnnxModel>,
}

impl FromResources for State {
    fn from_resources(resources: &Resources) -> Self {
        let asset_server = resources.get::<AssetServer>().unwrap();
        State {
            model: asset_server.load("model.onnx"),
        }
    }
}
```

And finally, I add my asset loader and my resource to the Bevy app:

```rust
App::build();
    .add_asset::<OnnxModel>()
    .init_asset_loader::<OnnxModelLoader>()
    .init_resource::<State>();
```

### Drawing to a texture
[To display a texture](https://github.com/vleue/bevmnist/blob/57b51cd5121f88af04e1defd40e9b8e1cb9c9623/src/main.rs#L166), I use standart Bevy UI components, here an `ImageBundle`. I also mark this entity with components `Interaction` and `FocusPolicy` from `bevy::ui` so that it can react to mouse clicks and movements.


To be input type independent (touch and mouse) I create an event `Draw` that takes a position coordinates. I can then trigger this event on `CursorMoved` events when mouse is clicked over the texture in a [system](https://github.com/vleue/bevmnist/blob/57b51cd5121f88af04e1defd40e9b8e1cb9c9623/src/main.rs#L249)

```rust
fn drawing_mouse(
    (mut reader, events): (Local<EventReader<CursorMoved>>, Res<Events<CursorMoved>>),
    mut last_mouse_position: Local<Option<Vec2>>,
    mut texture_events: ResMut<Events<Event>>,
    state: Res<State>,
    drawable: Query<(&Interaction, &GlobalTransform, &Style), With<Drawable>>,
) {
    for (interaction, transform, style) in drawable.iter() {
        if let Interaction::Clicked = interaction {
            // Get the width and height of the texture
            let width = if let Val::Px(x) = style.size.width {
                x
            } else {
                0.
            };
            let height = if let Val::Px(x) = style.size.height {
                x
            } else {
                0.
            };
            // For every `CursorMoved` event
            for event in reader.iter(&events) {
                if let Some(last_mouse_position) = *last_mouse_position {
                    // If mouvement is fast, interpolate positions between last known position
                    // and current position from event
                    let steps =
                        (last_mouse_position.distance(event.position) as u32 / INPUT_SIZE + 1) * 3;
                    for i in 0..steps {
                        let lerped =
                            last_mouse_position.lerp(event.position, i as f32 / steps as f32);

                        // Change cursor position from window to texture
                        let x = lerped.x - transform.translation.x + width / 2.;
                        let y = lerped.y - transform.translation.y + height / 2.;

                        // And send the event to draw at this position
                        texture_events.send(Event::Draw(Vec2::new(x, y)));
                    }
                } else {
                    let x = event.position.x - transform.translation.x + width / 2.;
                    let y = event.position.y - transform.translation.y + height / 2.;
                    texture_events.send(Event::Draw(Vec2::new(x, y)));
                }

                *last_mouse_position = Some(event.position);
            }
        } else {
            *last_mouse_position = None;
        }
    }
}
```

And to actually draw on the texture, I [listen for this event](https://github.com/vleue/bevmnist/blob/57b51cd5121f88af04e1defd40e9b8e1cb9c9623/src/main.rs#L338) and use a brush to color the texture around the event coordinates:

```rust
fn update_texture(
    (mut reader, events): (Local<EventReader<Event>>, Res<Events<Event>>),
    materials: Res<Assets<ColorMaterial>>,
    mut textures: ResMut<Assets<Texture>>,
    mut state: ResMut<State>,
    drawable: Query<(&bevy::ui::Node, &Handle<ColorMaterial>), With<Drawable>>,
) {
    for event in reader.iter(&events) {
        // Retrieving the texture data from it's `Handle`
        // First, getting the `Handle<ColorMaterial>` of the `ImageBundle`
        let (node, mat) = drawable.iter().next().unwrap();
        // Then, getting the `ColorMaterial` matching this handle
        let material = materials.get(mat).unwrap();
        // Finally, getting the texture itself from the `texture` field of the `ColorMaterial`
        let texture = textures
            .get_mut(material.texture.as_ref().unwrap())
            .unwrap();

        match event {
            Event::Draw(pos) => {
                // Use a large round brush instead of drawing pixel by pixel
                // `node.size` is the displayed size of the texture
                // `texture.size` is the actual size of the texture data
                // `INPUT_SIZE` is the expected input size by the model
                // The brush will be bigger if the drawing area is bigger to provide
                // a smoother drawing experience
                let radius = (1.3 * node.size.x / INPUT_SIZE as f32 / 2.) as i32;
                let scale = (texture.size.width as f32 / node.size.x) as i32;
                for i in -radius..(radius + 1) {
                    for j in -radius..(radius + 1) {
                        let target_point = Vec2::new(pos.x + i as f32, pos.y + j as f32);
                        if pos.distance(target_point) < radius as f32 {
                            for i in 0..=scale {
                                for j in 0..=scale {
                                    set_pixel(
                                        (target_point.x as i32) * scale + i,
                                        ((node.size.y as f32 - target_point.y) as i32) * scale + j,
                                        255,
                                        texture,
                                    )
                                }
                            }
                        }
                    }
                }
            }
        }
    }
}
```

### Getting model input from texture, and infering digit

I can now run my model on my texture and guess the digit!. This model is fast enough that this can run [at every frame](https://github.com/vleue/bevmnist/blob/57b51cd5121f88af04e1defd40e9b8e1cb9c9623/src/main.rs#L409)

```rust
fn infer(
    state: Res<State>,
    materials: Res<Assets<ColorMaterial>>,
    textures: Res<Assets<Texture>>,
    models: Res<Assets<OnnxModel>>,
    drawable: Query<&Handle<ColorMaterial>, With<Drawable>>,
    mut display: Query<&mut Text>,
) {
    for mat in drawable.iter() {
        // Get the texture from the `Handle<ColorMaterial>`
        let material = materials.get(mat).unwrap();
        let texture = textures.get(material.texture.as_ref().unwrap()).unwrap();

        // As the texture is much larger than the model input, each point in the
        // model input will be 1 if at least half of the point in a square
        // of `pixel_size`in the texture are colored
        let pixel_size = (texture.size.width as u32 / INPUT_SIZE) as i32;

        let image = tract_ndarray::Array4::from_shape_fn(
            (1, 1, INPUT_SIZE as usize, INPUT_SIZE as usize),
            |(_, _, y, x)| {
                let mut val = 0;
                for i in 0..pixel_size as i32 {
                    for j in 0..pixel_size as i32 {
                        val += get_pixel(
                            x as i32 * pixel_size + i,
                            y as i32 * pixel_size + j,
                            texture,
                        ) as i32;
                    }
                }
                if val > pixel_size * pixel_size / 2 {
                    1. as f32
                } else {
                    0. as f32
                }
            },
        )
        .into();

        if let Some(model) = models.get(state.model.as_weak::<OnnxModel>()) {
            // Run the model on the input
            let result = model.model.run(tvec!(image)).unwrap();

            // Get the best prediction, and display it if its score is high enough
            if let Some((value, score)) = result[0]
                .to_array_view::<f32>()
                .unwrap()
                .iter()
                .cloned()
                .enumerate()
                .max_by(|a, b| a.1.partial_cmp(&b.1).unwrap())
            {
                if score > 10. {
                    display.iter_mut().next().unwrap().value = format!("{:?}", value);
                } else {
                    display.iter_mut().next().unwrap().value = "".to_string();
                }
            }
        }
    }
}
```


## Build for Wasm target

Thanks to [bevy_webgl2](https://github.com/mrk-its/bevy_webgl2), this is actually very straightforward. I just need to [add the plugin `WebGL2Plugin`](https://github.com/vleue/bevmnist/blob/57b51cd5121f88af04e1defd40e9b8e1cb9c9623/src/main.rs#L84) and disable the default features of Bevy to only enable the one available on Wasm.

There is a [bevy webgl2 template](https://github.com/mrk-its/bevy_webgl2_app_template) if you want to build a game that can works on all platforms.


## Github Actions to deploy to [itch.io](https://vleue.itch.io)

To host this POC (and my other games) on [itch.io](https://vleue.itch.io), I have a workflow on Github Actions that build the game for Windows, Linux, macOS and Wasm, create a release on Github, and push everything to [itch.io](https://vleue.itch.io). This workflow is triggered on tags.

It is made of two jobs. [The first one](https://github.com/vleue/bevmnist/blob/57b51cd5121f88af04e1defd40e9b8e1cb9c9623/.github/workflows/release.yaml#L15) will, on each platform:
- setup the environment, install dependencies and tools
- build the game
- perform platform specific steps (`strip`, `wasm-bindgen`, ...) and copy assets if needed
- create an archive for the platform (`dmg` for macOS, `zip` for the other)
- create a release on Github and add those archives to the release
- save the archives as artifacts

[The second job](https://github.com/vleue/bevmnist/blob/57b51cd5121f88af04e1defd40e9b8e1cb9c9623/.github/workflows/release.yaml#L149) will take the artifacts from the first job and send them to [itch.io](https://vleue.itch.io) using butler, [itch.io official command line](https://itch.io/docs/butler/)