tfdeploy 0.0.4

Tiny, no-nonsense, self contained, TensorFlow inference
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

extern crate protoc_rust;

use std::{fs, io, path};

use std::io::{BufRead, Write};

fn main() {
    let inputs: Vec<path::PathBuf> = fs::read_dir("protos/tensorflow/core/framework")
        .unwrap()
        .map(|entry| entry.unwrap().path())
        .collect();

    protoc_rust::run(protoc_rust::Args {
        out_dir: &*std::env::var("OUT_DIR").unwrap(),
        input: &*inputs
            .iter()
            .map(|a| a.to_str().unwrap())
            .collect::<Vec<&str>>(),
        includes: &["protos"],
    }).expect("protoc");

    for input in inputs {
        let mut broken = path::PathBuf::from(std::env::var("OUT_DIR").unwrap())
            .join(input.file_name().unwrap());
        println!("cargo:rerun-if-changed={}", input.to_str().unwrap());
        let mut fixed = broken.clone();
        fixed.set_extension("rs");
        broken.set_extension("rs.orig");
        println!("rename {:?} {:?}", fixed, broken);
        fs::rename(&fixed, &broken).unwrap();
        let mut f = fs::File::create(fixed).unwrap();
        for line in io::BufReader::new(fs::File::open(broken).unwrap()).lines() {
            let line = line.unwrap();
            if !line.starts_with("#![") {
                writeln!(f, "{}", line).unwrap();
            }
        }
    }
}