Tensorflow Rust实战上篇

我的专栏地址：我的segmentfault,欢迎浏览

---

机器学习的一个方向是能够将它用于照片中的对象识别。这包括能够挑选动物，建筑物甚至人脸等特征。本文将引导您使用一些现有模型来使用rust和tensorflow完成人脸检测。我们将使用一个名为mtcnn的预训练模型进行人脸检测（注意：训练新模型不是我们在本文中关注的内容）。

挑战

我们想要读取照片，检测到人脸，然后返回带有绘制边界框的图像。
换句话说，我们想转换它（图片使用RustFest的许可，由FionaCastiñeira拍摄）：

Tensorflow and MTCNN

最初的MTCNN模型是使用Caffe编写的，但幸运的是mtcnn有许多tensorflow python实现。我将选择tf-mtcnn，因为它是直接转换为单个图像模型文件。

首先，我们要添加tensorflow rust作为依赖。从Cargo.toml开始：

[package]
name = "mtcnn"
version = "0.1.0"
edition = "2018"

[dependencies]
tensorflow = "0.12.0"

我们要做的是加载一个Graph，它是预先训练好的MTCNN，并运行一个Session。要点是Graph是用于计算的模型，Session是Graph的一次运行。有关这些概念的更多信息可以在这里找到。我喜欢将Graph视为大容器中的人造大脑，功能是输入图像和输出一些很棒的图像。

因此，让我们首先读取现有的mtcnn.pb模型并尝试加载它。 Tensorflow graph以protobuf格式序列化，可以使用Graph::import_graph_def加载。

use std::error::Error;

use tensorflow::Graph;
use tensorflow::ImportGraphDefOptions;

fn main() -> Result<(), Box<dyn Error>> {

    //First, we load up the graph as a byte array
    let model = include_bytes!("mtcnn.pb");

    //Then we create a tensorflow graph from the model
    let mut graph = Graph::new();
    graph.import_graph_def(&*model, &ImportGraphDefOptions::new())?

    Ok(())
}

跑cargo run命令，我们应该看到没有任何错误:

$ cargo run
   Compiling mtcnn v0.1.0 (~/mtcnn)
    Finished dev [unoptimized + debuginfo] target(s) in 0.89s
     Running `target/debug/mtcnn`

太棒了！看起来我们可以加载此图像！

处理命令行参数

StructOpt and the Command Line

我们想要测试图像生成，所以让我们使用structopt来获取两个参数：输入和输出。
如果您之前没有使用过structopt：structopt就像将clap与serde结合起来一样。输入参数是图像文件的路径。输出是我们保存输出图像的位置。

因此我们的结构体长这样：

use std::path::PathBuf;
use structopt::StructOpt

#[derive(StructOpt)]
struct Opt {
    #[structopt(parse(from_os_str))]
    input: PathBuf,

    #[structopt(parse(from_os_str))]
    output: PathBuf
}

parse（from_os_str）将字符串参数转换为PathBuf以节省一些开销，然后我们可以使用它来获取带有命令行参数的结构体：

fn main() -> Result<(), Box<dyn Error>> {

    let opt = Opt::from_args();
    ....
}

加载图像数据

我们需要提供以我们的图像数据为基础的tensorflow graph。那我们该怎么做呢？我们使用Tensor！Tensors表示我们图形中的数据，它让我想起将顶点发送到GPU。您有大量图像数据，并以tensorflow期望的格式发送给tensorflow。

此图中的输入tensor是一个浮点数组，其尺寸为：高x宽x 3（对于3个颜色通道）。

让我们使用image crate来加载图像：

let input_image = image::open(&opt.input)?;

接下来，我们要通过使用GenericImage::pixels函数将此图像转换为原始像素，并将其发送到我们的graph。所有多维Tensor阵列都是扁平的，并主要按行顺序存储。该模型使用BGR而不是传统的RGB颜色，因此我们需要在迭代时反转像素值。

把它们放在一起：

let mut flattened: Vec<f32> = Vec::new();

for (_x, _y, rgb) in input_image.pixels() {
    flattened.push(rgb[2] as f32);
    flattened.push(rgb[1] as f32);
    flattened.push(rgb[0] as f32);
}

这只是迭代像素，将它们添加到扁平的Vec。然后我们可以将它加载到tensor中，将图像高度和宽度指定为参数：

let input = Tensor::new(&[input_image.height() as u64, input_image.width() as u64, 3])
    .with_values(&flattened)?

太好了！我们已将图像加载到graph理解的格式。让我们跑一个session试试！

构建一个tensorflow session

我们已经准备好一个graph，我们可以使用默认值跑一个session:

let mut session = Session::new(&SessionOptions::new(), &graph)?;

跑一个session

在我们运行session之前，我们还有一些mtcnn模型所期望的其他输入参数。我们将使用与mtcnn库相同的默认值来获取这些值：

let min_size = Tensor::new(&[]).with_values(&[40f32])?;
let thresholds = Tensor::new(&[3]).with_values(&[0.6f32, 0.7f32, 0.7f32])?;
let factor = Tensor::new(&[]).with_values(&[0.709f32])?;

该graph可以定义在运行之前所需的多个输入/输出，并且它取决于具体的神经网络是什么。对于MTCNN，这些都在原始实现中描述。可能最重要的一个是min_size，它描述了寻找人脸的最小尺寸。

现在我们构建session的输入参数：

let mut args = SessionRunArgs::new();

//Load our parameters for the model
args.add_feed(&graph.operation_by_name_required("min_size")?, 0, &min_size);
args.add_feed(&graph.operation_by_name_required("thresholds")?, 0, &thresholds);
args.add_feed(&graph.operation_by_name_required("factor")?, 0, &factor);

//Load our input image
args.add_feed(&graph.operation_by_name_required("input")?, 0, &input);

好的，输出怎么样？session结束时我们要在结果里获取边界框和概率：

let bbox = args.request_fetch(&graph.operation_by_name_required("box")?, 0);
let prob = args.request_fetch(&graph.operation_by_name_required("prob")?, 0);

非常酷，我们已经定义好输入和输出，可以跑起来了！

session.run(&mut args)?;

输出处理

BBox结构体

该模型输出以下值：

• 边界框的人脸
• 人脸的地标
• 是人脸可能性：从0到1

为了使它更容易使用，我们将定义一个边界框结构，以更易于阅读的方式对这些值进行编码：

#[derive(Copy, Clone, Debug)]
pub struct BBox {
    pub x1: f32,
    pub y1: f32,
    pub x2: f32,
    pub y2: f32,
    pub prob: f32,
}

我们的工作是将我们从tensorflow session中返回的数组转换为此结构体，这样它就更有意义。

保存输出

是的，现在让我们保存输出。就像输入一样，输出也是Tensors：

let bbox_res: Tensor<f32> = args.fetch(bbox)?;
let prob_res: Tensor<f32> = args.fetch(prob)?;

bbox的定义是什么？好吧，它是一个多维扁平数组，每个边界框包含4个浮点数，代表边界框范围。 prob是一个浮点值的数组，每个人脸都有一个浮点值：从0到1的概率。所以我们应该期望bbox_res长度是人脸数x 4，而prob_res长度等于人脸数。

让我们做一些基本的迭代并将结果存储到Vec中：

//Let's store the results as a Vec<BBox>
let mut bboxes = Vec::new();

let mut i = 0;
let mut j = 0;

//While we have responses, iterate through
while i < bbox_res.len() {

    //Add in the 4 floats from the `bbox_res` array. 
    //Notice the y1, x1, etc.. is ordered differently to our struct definition.
    bboxes.push(BBox {
        y1: bbox_res[i],
        x1: bbox_res[i + 1],
        y2: bbox_res[i + 2],
        x2: bbox_res[i + 3],
        prob: prob_res[j], // Add in the facial probability
    });

    //Step `i` ahead by 4. 
    i += 4;
    //Step `i` ahead by 1. 
    j += 1;
}

打印边界框

好吧，我们还没有将边界框编码成图像。但是让我们进行debug，以确保我们得到了结果：

println!("BBox Length: {}, Bboxes:{:#?}", bboxes.len(), bboxes);

跑上面的代码，得到如下输出：

BBox Length: 120, BBoxes:[
    BBox {
        x1: 471.4591,
        y1: 287.59888,
        x2: 495.3053,
        y2: 317.25327,
        prob: 0.9999908
    },
    ....

哇哦!120张人脸!非常好！

绘制边界框

太好了，我们获得了一些边界框。让我们在图像上绘制它们，并保存到文件中。
要绘制边界框，我们可以使用imageproc库在边界框周围绘制简单边框。
首先，我们将在main函数之外定义线条颜色：

const LINE_COLOUR: Rgba<u8> = Rgba {
        data: [0, 255, 0, 0],
};

我们只读取输入图像，如果我们想要修改图像,首先要clone它：

let mut output_image = input_image.clone();

然后我们迭代bboxe数组:

for bbox in bboxes {
    //Drawing Happens Here!
}

接下来，我们使用draw_hollow_rect_mut函数。这将采用可变图像引用，并绘制由输入Rect指定的空心矩形（轮廓），覆盖任何现有像素。

Rect使用at函数获取x和y坐标，然后使用of_size函数获取宽度和高度。我们使用一些几何运算将我们的边界框转换为这种格式：

let rect = Rect::at(bbox.x1 as i32, bbox.y1 as i32)
    .of_size((bbox.x2 - bbox.x1) as u32, (bbox.y2 - bbox.y1) as u32);

然后绘制Rect:

draw_hollow_rect_mut(&mut img, rect, LINE_COLOUR);

当for循环结束，我们保存图像到输出文件：

output_image.save(&opt.output)?

ok,完成！我们跑一跑：

$ cargo run rustfest.jpg output.jpg
   Compiling mtcnn v0.1.0 (~/mtcnn)
    Finished dev [unoptimized + debuginfo] target(s) in 5.12s
     Running `target/debug/mtcnn rustfest.jpg output.jpg`
2019-03-28 16:15:48.194933: I tensorflow/core/platform/cpu_feature_guard.cc:141] Your CPU supports instructions that this TensorFlow binary was not compiled to use: SSE4.2 AVX AVX2 FMA
BBox Length: 154, BBoxes:[
    BBox {
        x1: 951.46875,
        y1: 274.00577,
        x2: 973.68304,
        y2: 301.93915,
        prob: 0.9999999
    },
....

太好了！没有错误！

总结

让我们逐步完成在这个小应用程序中所做的事情：

• 加载了预先训练的tensorflow graph
• 解析命令行参数
• 读入图像数据
• 通过运行tensorflow session提取人脸
• 将该session的结果绘制在原图像之上
• 将绘制结果保存到图像文件

希望这能为您提供在rust中使用tensorflow的良好介绍