Tensorflow Rust实战上篇

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

挑战

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

Tensorflow Rust实战上篇

Tensorflow Rust实战上篇

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就像将clapserde结合起来一样。输入参数是图像文件的路径。输出是我们保存输出图像的位置。

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

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的良好介绍

相关推荐