羚珑中有很多丰富的模板素材,人工对图片进行打标分类的话具有较主观的意识,且工作量也较为繁重。而利用深度学习技术对图片进行分类是当下一个很广泛的应用,于是我们考虑是否可以探索深度学习来对模板进行风格的分类识别。如下图所示是我们目前训练的一些预测效果,目前做到测试集 91%
的预测准确度:
数据集
上述风格预测使用的是卷积神经网络(Convolutional Neural Network 后面简称CNN
)对模板图片进行不同风格分类的识别。这个神经网络在一个已经做好分类的模板图片数据集上进行训练,该数据集包含 8 个分类的约 300 个模板图片。数据集虽然有点小,但是作为一个初步探索,问题还是不大的,而且随着后续数据的增加,训练精度也会越来越高。
原理流程
Google 开源的 TensorFlow 框架(后面简称 TF
),使我们轻而易举地快速编写机器学习模型并进行训练,而且还提供 JS 版本(TensorFlowJS 后面简称 TFJS
),可以直接在前端和服务端(NodeJS
)上使用,这对前端工程师来说实在是非常地友好。所以我们将会采用 TFJS
在 Node 服务器上进行风格分类的开发,开发流程如下图所示:
项目搭建
CNN
对图片的训练需要图片数据,而我们项目中的图片数据属于比较敏感的内容不方便对外开放,因此本文不会提供完整的项目实现代码,仅仅展示一些关键核心代码给大家了解。
我们直接使用 npm init
搭建一个 node
项目,然后我们需要用到这两个核心的库:
@tensorflow/tfjs-node
:TFJS
专门运行在NodeJS
上的库, 目前还处于开发阶段,还未稳定,安装该库可能需要科学上网,opencv4nodejs
:可以在NodeJS
上运行opencv
的库,主要要来读取图片元数据和处理压缩,安装该库可能有点点大 (- -! 巨大)。- 其他一些辅助的库就不介绍了。。。
可以通过 npm install
或者 yarn
的方式来安装。我们项目中的 package.json
如下所示:
{
...
"dependencies": {
"@tensorflow/tfjs": "^0.14.2",
"@tensorflow/tfjs-node": "^0.2.3",
"apisauce": "^1.0.2",
"blob": "0.0.5",
"get-pixels": "^3.3.2",
"jpeg-js": "^0.3.4",
"koa": "^2.6.2",
"koa-router": "^7.4.0",
"lodash": "^4.17.11",
"nodemon": "^1.18.9",
"opencv4nodejs": "^4.14.0"
}
}
数据处理
1.图片归一化处理
我们平时遇到的大多数 JPG
、PNG
图片都是位图,也叫点阵图,像素图,对计算机来说其本质是一个多维矩阵 。而图像的分析和处理其实是对这个多维矩阵的一种运算,理解这一点是我们整个处理流程的核心。同时图像的色彩模式也有很多种,例如 RGB 、CMYK、HSV 等等,我们这个项目讨论的是 RGB 模式位图来解决运算的问题。 RGB 模式的图片包含 3 个通道分别对应 Red
、Green
、Blue
3 种颜色,每个通道分别包含着图像上每个像素在对应该通道颜色的数值,所以 RGB 图片对我们来说其实是一个 3 维的矩阵数组。
CNN
对图片进行训练,而 CNN
训练时需要将图片归一化成统一的尺寸,所以我们需要对每张图片的尺寸进行特定压缩才可以进行训练和预测。这里我们为了尽快看到效果,和减少计算量,先忽略图片质量,暂时将每一张图片统一压缩成 32 x 32
的尺寸。
PS: 这里的压缩会是一个有争议的地方,因为压缩后会导致图片变形,有其他更加优化的算法可以解决,但这是一个很大的话题,本文先不做讨论
其次为了把图片色彩放到考虑的范围内,我们保留图片的色彩信息,也就是保留 3 个通道的信息,这样我们最后会得到一个形状为 [32, 32, 3]
的矩阵数组。
我们将会按照标记好分类的图片目录,把所有的图片处理一遍。主要的代码我们会用到 opencv4nodejs
当中的 resize
方法来完成图片尺寸的压缩,如下:
const tf = require("@tensorflow/tfjs");
const fs = require("fs");
const cv = require("opencv4nodejs");
const { forEach, startsWith } = require("lodash");
const ImageWidth = 32;
const ImageHeight = 32;
function loadImgFromDir(dir_path) {
const dirs = fs.readdirSync(dir_path);
const datas = [];
forEach(dirs, dirName => {
const dir = fs.lstatSync(dir_path + dirName);
if (!dir.isDirectory()) return; // 非目录的文件跳过
const files = fs.readdirSync(dir_path + dirName);
forEach(files, fileName => {
if (startsWith(fileName, ".")) return; // 如果以 .开头则跳过文件
const path = dir_path + dirName + "/" + fileName;
const data = fs.readFileSync(path);
const image = cv.imread(path);
const res = image.resize(ImageWidth, ImageHeight);
datas.push({ image: res.getDataAsArray(), label: dirName });
});
});
return datas;
}
2.数据集拆分
机器学习包含监督式学习、无监督学习、增强学习(强化学习)等,我们的风格分类模型是一个典型的 监督式学习
。
在 监督式学习
中,数据集的通常划分为 2 ~ 3 个部分:
训练集
,学习样本数据集,通过匹配一些参数来建立一个分类器。建立一种分类的方式,主要是用来训练模型的。验证集
,对学习出来的模型,微调分类器的参数,如在神经网络中选择隐藏单元数。验证集还用来确定网络结构或者控制模型复杂程度的参数。测试集
,主要用于测试训练好的模型的分类能力
严格上来说,我们需要对数据集按照上述划分,各个数据集之间互不交叉,再进行 交叉验证 才能得到一个比较好的分类效果,但在实际应用中,一般可以将数据集随机打乱,再按照经验以 3:1 的比例划分成训练集和测试集。将训练集作为训练样本,然后用测试集描述模型的泛化能力,把测试集当成一个简单的交叉验证集使用。
划分数据集的时候我们一般会顺便把训练集转化成 TF
中的 Tensor(张量)
,其本质是一个矩阵,以方便后面对模型进行训练。然后还会读取每张图片标记好的分类,转化成 one-hot-code(独热编码),按顺序和图片一一对应记录。核心代码如下:
const tf = require('@tensorflow/tfjs')
const { forEach, shuffle } = require('lodash')
function loadData() {
const datas = loadImgFromDir(DIR_PATH)
// 打乱顺序
const shuffleDatas = shuffle(datas)
const threshold = 0.75 * shuffleDatas.length
const trainImages = []
const trainLabels = []
const testImages = []
const testLabels = []
forEach(shuffleDatas, (data, index) => {
if (index < threshold) {
// 小于阈值的都认为训练集
trainImages.push(data.image)
trainLabels.push(data.label)
} else {
// 大于阈值的都认为测试集
testImages.push(data.image)
testLabels.push(data.label)
}
})
let b = new Buffer(shuffleDatas)
return {
trainImages: tf.tensor4d(trainImages),
trainLabels: tf.oneHot(tf.tensor1d(trainLabels, 'int32'), 8).toFloat(),
testImages: tf.tensor4d(testImages),
testLabels: tf.oneHot(tf.tensor1d(testLabels, 'int32'), 8).toFloat(),
}
}
模型训练
1. 模型搭建
我们将搭建一个经典的 CNN
模型来进行训练,该模型包含 2 个卷积层和 2 个全连接层,输入是一个形状为 [32, 32, 3]
的图片矩阵数组,输出是一个表示各个分类概率的浮点型 1 维数组,学习速率设置为 0.001, 由于是分类问题,我们将采用 categoricalCrossentropy
(多类的损失函数,针对分类模型优化的交叉嫡损失函数) 优化器。这是一个相对简单但经典的 CNN
模型,模型越简单,参数越少,训练速度也会越快,模型越复杂,参数就越多,训练速度会比较慢,相对来说效果一般会更好,但对训练的机器要求就会比较高。我们的机器是在 docker
上运行的,所以选择了一个简单的模型先,模型代码如下:
const tf = require('@tensorflow/tfjs')
const LEARNING_RATE = 0.001
const model = tf.sequential()
// 第1层卷积网络
model.add(tf.layers.conv2d({
inputShape: [32, 32, 3],
filters: 32,
kernelSize: 3,
activation: 'relu',
}))
// 最大池化
model.add(tf.layers.maxPooling2d({
poolSize: [2, 2]
}))
// 第2层卷积网络
model.add(tf.layers.conv2d({
filters: 32,
kernelSize: 3,
activation: 'relu',
}))
// 最大池化
model.add(tf.layers.maxPooling2d({
poolSize: [2, 2]
}))
// 拉成1维数组
model.add(tf.layers.flatten())
// 随机失活
model.add(tf.layers.dropout({rate: 0.35}))
// 全连接层
model.add(tf.layers.dense({ units: 128, activation: 'relu'}))
// 随机失活
model.add(tf.layers.dropout({rate: 0.5}))
// 全连接层2
model.add(tf.layers.dense({units: 8, activation: 'softmax'}))
const optimizer = tf.train.adam(LEARNING_RATE)
model.compile({optimizer, loss: 'categoricalCrossentropy', metrics: ['accuracy']})
module.exports = model
2.模型训练
创建完模型之后我们就可以开始进行训练了,先读出模型后 compile
一下,然后输入训练集图片数组和对应的标签数组,再执行 model.fit()
并设置参数 epochs (训练次数)
为 50 ,设置 batchSize (每个训练批次的最大样本数)
为 128, 就可以直接进行训练了。训练完之后我们输入之前拆分出来的测试集数据进行一下测试,看看准确率能达到多少,以衡量该模型的泛化能力。主要代码如下:
const tf = require('@tensorflow/tfjs')
require('@tensorflow/tfjs-node');
const model = require('./model')
async function train () {
const { trainImages, trainLabels, testImages, testLabels } = await loadData()
await model.fit(trainImages, trainLabels, {
epochs: 50,
batchSize: 128,
})
console.log('evluate Test...')
const evalOutput = model.evaluate(testImages, testLabels)
console.log(
`\nEvaluation result:\n` +
` Loss = ${evalOutput[0].dataSync()[0].toFixed(3)}; `+
`Accuracy = ${evalOutput[1].dataSync()[0].toFixed(3)}`)
...
}
训练过程入下图所示:
从训练过程的日志我们可以看出来,经过 50 轮的训练之后,我们的训练集预测准确精度达到 97.7%
,而在测试集的预测准确精度达到了 91.7%
,这个结果说明模型的泛化能力还是不错的。
训练完成之后我们需要对模型进行保存,以方便我们后续进行使用,可以再上面训练完成后的地方插入 model.save()
语句并指定一个目录 model
进行保存。
async function train () {
...
await model.save('file://model')
}
保存完成后我们将得到一个 model
目录,里面包含 model.json
和 weights.bin
两个文件,他们保存了我们这个风格预测模型的结构和训练后的权重信息,得到这个模型之后可以将它转化成我们想要运行的任何平台(iOS、Android、linux 设备、网页等等)、任何语言(JS、Python、Swift、C/C++、Java 等等)所支持的模型文件,可以随处使用!(^ _ ^ cool! )
风格预测
1. 读取模型
在 NodeJS
中,我们可以从前面步骤保存的模型文件目录直接通过 tf.loadModel()
方法读取模型,非常地方便!在读取模型之后,我们并不能直接使用,还需要使用 compile
方法对模型设置一次跟之前相同的学习速率和优化函数。
async function predict () {
const model = await tf.loadModel('file://model/model.json')
const optimizer = tf.train.adam(LEARNING_RATE)
model.compile({optimizer, loss: 'categoricalCrossentropy', metrics: ['accuracy']})
...
}
2.模型预测
读取完模型之后,我们将需要预测的图片再次归一化处理成训练时输入的形状 [32, 32, 3]
,并转化成 Tensor
类型,再调用 TFJS
提供的 model.predict()
方法即可得到预测概率的数组,再结合风格分类类型,简单地取预测概率最高的前 3 项,得到预测结果的一个判断。核心代码如下:
// 从 RGBA 元数据数组获取图片 tensor
function getTensorFromPixels(pixels) {
const numChannels = 3
const numPixels = pixels.width*pixels.height
const values = new Uint8Array(numPixels * numChannels)
for (let i = 0; i < numPixels; i++) {
for (let channel = 0; channel < numChannels; ++channel) {
values[i * numChannels + channel] = pixels.data[i * 4 + channel]
}
}
return tf.tensor3d(values, [pixels.width, pixels.height, numChannels])
}
// 根据预测概率结果和风格类型计算概率最高的前 3 项预测结果
function getResultFromPredict (predict) {
const result = map(predict, (p, i) => {
return {
probability: p,
className: CLASS_NAME[i],
}
}).sort((a, b) => {
return b.probability - a.probability
})
return slice(result, 0, 3)
}
async function predict (url) {
...
// 预测图片
const pixels = await getPixels(url) // 从图片地址获取图片二进制内容,方法有很多,此处代码省略
const imgTensor = getTensorFromPixels(pixels)
const normalizedImg = tf.image.resizeNearestNeighbor(imgTensor, [32, 32]).reshape([1, 32, 32, 3])
const predict = await model.predict(normalizedImg).data()
const result = getResultFromPredict(predict)
}
最终我们输入需要预测的图片地址后可以得到如下图所示的预测结果:
后续展望
目前风格识别还是一个比较简单的单分类模型,我们会继续改进成一个多分类的模型。我们还希望打通羚珑目前的模板数据库,抽离出一个模板数据池,通过一些自动化的流程来持续改进模型算法,提升预测准确率和泛化能力。
另外,目前对深度学习的探索还处于一个初步的阶段,羚珑中还有更多的智能场景等待着我们不断地去挖掘和探索,让羚珑更加智能地助力各个业务的设计场景需求。