cervo 0.2.0

High level APIs for RL in games.
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

// Author: Tom Solberg <tom.solberg@embark-studios.com>
// Copyright © 2022, Embark Studios AB, all rights reserved.
// Created: 13 May 2022

/*!

*/

use anyhow::{bail, Result};
use cervo_asset::{AssetData, AssetKind};
use clap::Parser;
use std::path::PathBuf;

/// Package the file into a crvo file.
#[derive(Parser, Debug)]
#[clap()]
pub(crate) struct PackageArgs {
    /// The file to load. Must be .onnx or .nnef.tar
    infile: PathBuf,
    /// The file to write. Extensions are ignored.
    outfile: PathBuf,

    /// If set and reading an ONNX file, convert to NNEF on packaging.
    #[clap(short = 'O', long = "optimize")]
    optimize: bool,

    /// If provided, specialize the model for the batch size.
    #[clap(short = 'N', long = "batch-size", requires = "optimize")]
    batch_size: Option<usize>,
}

pub(super) fn package(config: PackageArgs) -> Result<()> {
    let kind = if cervo_nnef::is_nnef_tar(&config.infile) {
        AssetKind::Nnef
    } else {
        match config.infile.extension().and_then(|ext| ext.to_str()) {
            Some("onnx") => AssetKind::Onnx,
            Some(other) => bail!("unknown file type {:?}", other),
            None => bail!("missing file extension {:?}", config.infile),
        }
    };

    let bytes = std::fs::read(&config.infile)?;
    let asset = AssetData::new(kind, bytes);

    let asset = match asset.kind() {
        AssetKind::Onnx => {
            if config.optimize {
                asset.to_nnef(config.batch_size)?
            } else {
                asset
            }
        }
        AssetKind::Nnef => {
            if config.optimize {
                bail!("cannot optimize NNEF assets further")
            } else {
                asset
            }
        }
    };

    let file = config.outfile.with_extension("crvo");
    let data = asset.serialize()?;
    std::fs::write(&file, data)?;
    Ok(())
}