Java调用TensorFlow库预测图片质量

概述

本文参考官方的java使用TensorFlow库的例子，将深度学习模型保存成pb文件，在java环境中加载模型并做预测。

环境安装

安装pip

yum -y install epel-release
yum install python-pip
pip install --upgrade pip

安装TensorFlow、Keras、numpy

pip install tensorflow  //安装的是最新的tensorflow2.1版本
pip install keras
pip install numpy

Maven配置

在pom.xml中增加如下配置，加载java的tensorflow库

<dependency>
    <groupId>org.tensorflow</groupId>
    <artifactId>tensorflow</artifactId>
    <version>1.15.0</version>
</dependency>

加载模型

InputStream inputStream = ImageRecognize.class.getResourceAsStream(MODEL_PATH);
Graph graph = new Graph();
graph.importGraphDef(IOUtils.toByteArray(inputStream));
Session session = new Session(graph);

其中graph和session都是线程安全的，可以使用单例，不需要每次请求都重新加载模型和Session。

图片预处理

我们使用的是Xception模型，它要求输入的图片大小是480 * 480 * 3，并且需要对图片做预处理，每个rgb值归一化到[-1,1]区间。下面我们介绍两种预处理的方式：

BufferedImage预处理

BufferedImage bufferedImage = new BufferedImage(480, 480, BufferedImage.TYPE_INT_RGB);
Graphics graphics = bufferedImage.getGraphics();

InputStream in = new ByteArrayInputStream(imageData);
Image srcImage = ImageIO.read(in);
graphics.drawImage(srcImage, 0, 0, 480, 480, null); //将图片大小转换为480*480

int w = bufferedImage.getWidth();
int h = bufferedImage.getHeight();
float[][][][] imgTensor = new float[1][h][w][3];
for (int i = 0; i < h; i++) {
     for (int j = 0; j < w; j++) {
              int pixel = bufferedImage.getRGB(j, i); // 下面三行代码将一个数字转换为RGB数字，同时归一化到[-1,1]区间
              imgTensor[0][i][j][0] = (float) ((pixel & 0xff0000) >> 16) / 127.5f - 1;
              imgTensor[0][i][j][1] = (float) ((pixel & 0xff00) >> 8) / 127.5f - 1;
              imgTensor[0][i][j][2] = (float) ((pixel & 0xff)) / 127.5f - 1;
      }
 }
return Tensors.create(imgTensor);

TensorFlow预处理

TensorFlow的预处理参考了LabelImage.java调用方式，它是使用TensorFlow Graph的一些预定义好的Operator来对图片做预处理。

  private Tensor<Float> getImageTensor(byte[] imageBytes){
      Graph g = new Graph();
      GraphBuilder b = new GraphBuilder(g);

      final int H = IMAGE_HEIGTH;
      final int W = IMAGE_WIDTH;
      final float mean = 1f;
      final float scale = 127.5f;

      final Output<String> input = b.constant("input", imageBytes);
      final Output<Float> output =
              b.sub(
                      b.div(
                              b.resizeBilinear(
                                      b.expandDims(
                                              b.cast(b.decodeJpeg(input, 3), Float.class), //解析jpeg文件
                                              b.constant("make_batch", 0) //扩展成4维Tensor
                                      ),
                                      b.constant("size", new int[]{H, W}) //resize图片成[H,W]大小
                              ),
                              b.constant("scale", scale) //每个值除以127.5f
                      ),
                      b.constant("mean", mean) //归一化到[-1,1]区间
              );
      try (Session s = new Session(g)) {
        // Generally, there may be multiple output tensors, all of them must be closed to prevent resource leaks.
        return s.runner().fetch(output.op().name()).run().get(0).expect(Float.class);
      }
  }

  static class GraphBuilder {
    GraphBuilder(Graph g) {
      this.g = g;
    }

    Output<Float> div(Output<Float> x, Output<Float> y) {
      return binaryOp("Div", x, y);
    }

    <T> Output<T> sub(Output<T> x, Output<T> y) {
      return binaryOp("Sub", x, y);
    }

    <T> Output<Float> resizeBilinear(Output<T> images, Output<Integer> size) {
      return binaryOp3("ResizeBilinear", images, size);
    }

    <T> Output<T> expandDims(Output<T> input, Output<Integer> dim) {
      return binaryOp3("ExpandDims", input, dim);
    }

    <T, U> Output<U> cast(Output<T> value, Class<U> type) {
      DataType dtype = DataType.fromClass(type);
      return g.opBuilder("Cast", "Cast")
              .addInput(value)
              .setAttr("DstT", dtype)
              .build()
              .<U>output(0);
    }

    Output<UInt8> decodeJpeg(Output<String> contents, long channels) {
      return g.opBuilder("DecodeJpeg", "DecodeJpeg")
              .addInput(contents)
              .setAttr("channels", channels)
              .build()
              .<UInt8>output(0);
    }

    <T> Output<T> constant(String name, Object value, Class<T> type) {
      try (Tensor<T> t = Tensor.<T>create(value, type)) {
        return g.opBuilder("Const", name)
                .setAttr("dtype", DataType.fromClass(type))
                .setAttr("value", t)
                .build()
                .<T>output(0);
      }
    }
    Output<String> constant(String name, byte[] value) {
      return this.constant(name, value, String.class);
    }

    Output<Integer> constant(String name, int value) {
      return this.constant(name, value, Integer.class);
    }

    Output<Integer> constant(String name, int[] value) {
      return this.constant(name, value, Integer.class);
    }

    Output<Float> constant(String name, float value) {
      return this.constant(name, value, Float.class);
    }

    private <T> Output<T> binaryOp(String type, Output<T> in1, Output<T> in2) {
      return g.opBuilder(type, type).addInput(in1).addInput(in2).build().<T>output(0);
    }

    private <T, U, V> Output<T> binaryOp3(String type, Output<U> in1, Output<V> in2) {
      return g.opBuilder(type, type).addInput(in1).addInput(in2).build().<T>output(0);
    }
    private Graph g;
  }

上面两种方式都做了尝试，我们是原始图片大小为640 * 640，机器是2.4GHz的CPU机器（没有用GPU），第一种预处理方法平均耗时在200ms左右，第二种预处理方法耗时为30ms左右，主要原因为TensorFlow内部对矩阵运算会做优化，而第一种方法手写的循环效率不高。后续会再测一下在GPU环境下的耗时。

模型预测

我们的xception模型中，输入节点的名字为input_1，输出节点的名字为output，对应着代码里的名字，需要完全一致。

    float result = -1;
    input = getImageTensor1(imageData);
    if ( input == null ) {
      return result;
    }

    List<Tensor<?>> results = session.runner().feed("input_1", input).fetch("output").run();
    if (results.size() > 0 && results.get(0).shape().length == 2) {
      long[] rshape = results.get(0).shape();
      int rs = (int) rshape[0];
      int rt = (int) rshape[1];
      float realResult[][] = new float[rs][rt];

      results.get(0).copyTo(realResult);
      for (int i = 0; i < rs; i++) {
        for (int j = 0; j < rt; j++) {
          result = realResult[i][j];
          break;
        }
      }
    }

线上部署

线上使用时候，有一个线程不断的从HDFS中检查并读取最新的模型。一旦模型有更新，则加载新模型替换旧模型。