如何利用TensorFlow.js部署简单的AI版“你画我猜”图像识别应用

选自Medium

作者:Zaid Alyafeai

机器之心编译

参与:Geek AI、路

本文创建了一个简单的工具来识别手绘图像,并且输出当前图像的名称。该应用无需安装任何额外的插件,可直接在浏览器上运行。作者使用谷歌 Colab 来训练模型,并使用 TensorFlow.js 将它部署到浏览器上。

代码和 demo

demo 地址:https://zaidalyafeai.github.io/sketcher/

代码地址:https://github.com/zaidalyafeai/zaidalyafeai.github.io/tree/master/sketcher

请通过以下链接在谷歌 Colab 上测试自己的 notebook:https://colab.research.google.com/github/zaidalyafeai/zaidalyafeai.github.io/blob/master/sketcher/Sketcher.ipynb

数据集

我们将使用卷积神经网络(CNN)来识别不同类型的手绘图像。这个卷积神经网络将在 Quick Draw 数据集(https://github.com/googlecreativelab/quickdraw-dataset)上接受训练。该数据集包含 345 个类别的大约 5 千万张手绘图像。

部分图像类别

流程

我们将使用 Keras 框架在谷歌 Colab 免费提供的 GPU 上训练模型,然后使用 TensorFlow.js 直接在浏览器上运行模型。我在 TensorFlow.js 上创建了一个教程(https://medium.com/tensorflow/a-gentle-introduction-to-tensorflow-js-dba2e5257702)。在继续下面的工作之前,请务必先阅读一下这个教程。下图为该项目的处理流程:

流程

在 Colab 上进行训练

谷歌 Colab 为我们提供了免费的 GPU 处理能力。你可以阅读下面的教程(https://medium.com/deep-learning-turkey/google-colab-free-gpu-tutorial-e113627b9f5d)了解如何创建 notebook 和开始进行 GPU 编程。

导入

我们将使用以 TensorFlow 作为后端、Keras 作为前端的编程框架

加载数据

由于内存容量有限,我们不会使用所有类别的图像进行训练。我们仅使用数据集中的 100 个类别(https://raw.githubusercontent.com/zaidalyafeai/zaidalyafeai.github.io/master/sketcher/mini_classes.txt)。每个类别的数据可以在谷歌 Colab(https://console.cloud.google.com/storage/browser/quickdrawdataset/full/numpybitmap?pli=1)上以 NumPy 数组的形式获得,数组的大小为 [N, 784],其中 N 为某类图像的数量。我们首先下载这个数据集:

由于内存限制,我们在这里将每类图像仅仅加载 5000 张。我们还将留出其中的 20% 作为测试数据。

数据预处理

我们对数据进行预处理操作,为训练模型做准备。该模型将使用规模为 [N, 28, 28, 1] 的批处理,并且输出规模为 [N, 100] 的概率。

创建模型

我们将创建一个简单的卷积神经网络。请注意,模型越简单、参数越少越好。实际上,我们将把模型转换到浏览器上然后再运行,并希望模型能在预测任务中快速运行。下面的模型包含 3 个卷积层和 2 个全连接层:

拟合、验证及测试

在这之后我们对模型进行了 5 轮训练,将训练数据分成了 256 批输入模型,并且分离出 10% 作为验证集。

训练结果如下图所示:

测试准确率达到了 92.20% 的 top 5 准确率。

准备 WEB 格式的模型

在我们得到满意的模型准确率后,我们将模型保存下来,以便进行下一步的转换。

为转换安装 tensorflow.js:

接着我们对模型进行转换:

这个步骤将创建一些权重文件和包含模型架构的 json 文件。

通过 zip 将模型进行压缩,以便将其下载到本地机器上:

最后下载模型:

在浏览器上进行推断

本节中,我们将展示如何加载模型并且进行推断。假设我们有一个尺寸为 300*300 的画布。在这里,我们不会详细介绍函数接口,而是将重点放在 TensorFlow.js 的部分。

加载模型

为了使用 TensorFlow.js,我们首先使用下面的脚本:

你的本地机器上需要有一台运行中的服务器来托管权重文件。你可以在 GitHub 上创建一个 apache 服务器或者托管网页,就像我在我的项目中所做的那样(https://github.com/zaidalyafeai/zaidalyafeai.github.io/tree/master/sketcher)。

接着,通过下面的代码将模型加载到浏览器:

关键字 await 的意思是等待模型被浏览器加载。

预处理

在进行预测前,我们需要对数据进行预处理。首先从画布中获取图像数据:

文章稍后将介绍 getMinBox()。dpi 变量被用于根据屏幕像素的密度对裁剪出的画布进行拉伸。

我们将画布当前的图像数据转化为一个张量,调整大小并进行归一化处理:

我们使用 model.predict 进行预测,这将返回一个规模为「N, 100」的概率。

我们可以使用简单的函数找到 top 5 概率。

提升准确率

请记住,我们的模型接受的输入数据是规模为 [N, 28, 28, 1] 的张量。我们绘图画布的尺寸为 300*300,这可能是两个手绘图像的大小,或者用户可以在上面绘制一个小图像。最好只裁剪包含当前手绘图像的方框。为了做到这一点,我们通过找到左上方和右下方的点来提取围绕图像的最小边界框。

用手绘图像进行测试

下图显示了一些第一次绘制的图像以及准确率最高的类别。所有的手绘图像都是我用鼠标画的,用笔绘制的话应该会得到更高的准确率。

本文为机器之心编译,转载请联系本公众号获得授权。

------------------------------------------------

  • 发表于:
  • 原文链接https://kuaibao.qq.com/s/20180722A0MPQ000?refer=cp_1026
  • 腾讯「云+社区」是腾讯内容开放平台帐号(企鹅号)传播渠道之一,根据《腾讯内容开放平台服务协议》转载发布内容。
  • 如有侵权,请联系 yunjia_community@tencent.com 删除。

扫码关注云+社区

领取腾讯云代金券