阅读 4503

手把手教你打造前端智能图标识别工具

Hi~ 我是前端学徒业枫(@Malpor),今天为大家带来一篇硬核前端智能化教程,真·手把手教你用机器学习打造一个纯前端运行的图标智能识别工具。并附上完整代码,一起来体验前端智能化的魅力吧~

背景

目前的前端组件库都使用 Iconfont 来管理图标,随着时间推移,图标越来越多,图标的命名也五花八门,很难约束。开发者还原设计稿时,经常要人肉从几百个图标中寻找对应的图标。有时候连设计师都找不到,导致重复添加图标。

最近发现在 AntDesign 官网有以图搜图标的功能,用户对设计稿或任意图片中的图标截图,点击/拖拽/粘贴上传,就可以搜索到匹配度最高的几个图标:AntDesign Icon功能开发者文章

v2-fee8f94fbfe1885a713cd6fbb854385a_b.gif

这个功能很好的解决了上面提到的问题,但还有些不足:

  • 截图最好是正方形的,否则拉伸后识别率会下降(后面会解释)。
  • 只能识别 AntDesign 的图标。

为了解决这些问题,我们决定自己打造一个前端图标识别工具。下面将以我们团队的开源组件库 Cloud Design 为例,手把手教你打造纯前端的专属图标识别工具。(完整代码放在文末)

术语简介

简单介绍几个术语,了解的同学可以直接跳过。

机器学习

机器学习研究和构建的是一种特殊算法(而非某一个特定的算法),能够让计算机自己在数据中学习从而进行预测。

所以,机器学习不是某种具体的算法,而是很多算法的统称。

机器学习包含:线性回归、贝叶斯、聚类、决策树、深度学习等等。前面 AntDesign 的模型是通过深度学习的代表算法 CNN 训练得到的。

CNN 卷积神经网络

卷积神经网络(Convolutional Neural Networks, CNN)是一类包含卷积计算且具有深度结构的前馈神经网络(Feedforward Neural Networks),最常用于分析视觉图像。

CNN 能有效的将大数据量图片降维到小数据量,且保留图像特征,非常适合处理图像数据。即使图像翻转、旋转或变换位置也能有效识别,常用来解决:图像分类检索、目标定位监测、人脸识别等等。

开始行动吧

我们要对图标进行识别,属于机器学习中经典的“图像分类”问题。CNN(卷积神经网络) 可以有效的识别图标,但是无法适应拉伸变形的场景。因为模型输入时要先把图像变换为正方形尺寸,截图尺寸不对会导致图像拉伸变形,降低识别率,甚至识别错误。

截屏2021-03-21 下午10.36.41.png

常用的解法有两种:

1、纯机器学习:通过增加不同拉伸状态的样本,让模型适应变形的图像。

2、机器学习 + 图像处理:用图像处理算法对数据进行裁剪,保证图像接近正方形。

第一种方法需要生成大量的训练数据,训练速度变慢,而且拉伸变形的情况很难遍历。第二种方法只需要进行简单的图像处理就可以有效提高识别率,所以我选择了它。那最终工作流应该是这样的:

截屏2021-03-21 下午10.38.04.png

接下来我会从 样本生成、模型训练、模型使用 三部分来介绍完整的过程。

样本生成

图像分类的训练样本都是图片,我们的图标则是 iconfont 渲染在页面中的。可以自然想到用 样本页面 + Puppeteer 截图来生成样本。但截图速度很慢,我也不想用 Faas 服务,于是想了个本地生成的方法:

首先人工把图标库的css部分转为js:

截屏2021-03-21 下午10.40.28.png

这样就能把图标当作文本绘制在 canvas 上,并用图像算法裁剪四周的空白区域:

// 用离屏 canvas 绘制图标
offscreenCtx.font = `20px NextIcon`;
offscreenCtx.fillText(labelMap[labelName]);

// 用 getImageData 获取图片数据,计算需裁剪的坐标
const { x, y, width: w, height: h } = getCutPosition(canvasSize, canvasSize, offscreenCtx.getImageData(0, 0, canvasSize, canvasSize).data);

// 计算需裁剪的坐标
function getCutPosition(width, height, imgData) {
  let lOffset = width; let rOffset = 0; let tOffset = height; let bOffset = 0;
  // 遍历像素,获取最小的非空白矩形区域
  for (let i = 0; i < width; i++) {
    for (let j = 0; j < height; j++) {
      const pos = (i + width * j) * 4;
      if (notEmpty(imgData[pos], imgData[pos + 1], imgData[pos + 2], imgData[pos + 3])) {
        // 调整 lOffset、rOffset、tOffset、bOffset
        // 略
      }
    }
  }
  
  // 如果形状不是正方形,将其扩展为正方形
  const r = (rOffset - lOffset) / (bOffset - tOffset);
  if (r !== 1) {
    // 略
  }

  return { x: lOffset, y: tOffset, width: rOffset - lOffset, height: bOffset - tOffset };
}

// 阈值 0 - 255
const d = 5;
// 判断是否非空白像素
function notEmpty(r, g, b, a) {
  return r < 255 - d && g < 255 - d && b < 255 - d;
}

// 用 canvas 裁剪 & 缩放图像,导出为 base64
ctx.drawImage(offscreenCanvas, x, y, w, h, 0, 0, 96, 96);
canvas.toDataURL('image/jpeg');
复制代码

截屏2021-03-21 下午10.42.59.png

生成一张图片的逻辑就写完了。改造一下,遍历不同图标、不同字号,可以得到全量的样本:

const fontStep = 1;
const fontSize = [20, 96];

labels.map((labelName) => {
  // 遍历不同的字号绘制图标
  for (let i = fontSize[0]; i <= fontSize[1]; i += fontStep) {
    // ...before
    offscreenCtx.font = `${i}px NextIcon`;
    // 其它逻辑
  }
});
复制代码

通过 Blob 将数据作为一个 json 下载:

const resultData = /* 生成全量数据 */;

const aLink = document.createElement('a');
const blob = new Blob([JSON.stringify(resultData, null, 2)], { type : 'application/json' });
aLink.download = 'icon.json';
aLink.href = URL.createObjectURL(blob);
aLink.click();
复制代码

这样就得到了包含几万张(350个图标,每个分类约70张图)样本图片的大 json,大概长这样:

[
  {
    "name": "smile",
    "data": [
      {
        "url": "data:image/jpeg;base64,/9j/4AA...IkB//9k=",
        "size": 20
      },
      {
        "url": "data:image/jpeg;base64,/9j/4AA...JAf//Z",
        "size": 21
      },
      ...
    ]
  },
]
复制代码

最后写一个简单的 node 程序,把每个分类的样本按照训练集70%,验证集20%,测试集10%的比例拆分打散并存储为图片文件。

--- train
  |-- smile
    |-- smile_3.jpg
    |-- smile_7.jpg
  |-- cry
    |-- cry_2.jpg
    |-- cry_8.jpg
  ...
--- validation
  |-- smile
  |-- cry
  ...
--- test
  |-- smile
  |-- cry
  ...
复制代码

这样我们就得到了完整的训练样本,而且生成速度很快,运行一遍只要1分钟左右。然后把三个目录一起打包成一个 zip 文件即可,因为下一步训练只支持 zip 格式。

模型训练

机器学习工具有很多种,作为一个前端,我最终选择使用 Pipcook 来训练。

Pipcook 项目是一个开源工具集,它能让 Web 开发者更好地使用机器学习,从而开启和加速前端智能化时代!

Pipcook 的安装和教程看官网(链接)即可,要注意目前只支持 Mac & Linux,Windows 暂时无法使用(Windows 可以使用 Tensorflow.js 训练)。

写一份 pipcook 的配置项:

{
  "plugins": {
    "dataCollect": {
      "package": "@pipcook/plugins-image-classification-data-collect",
      "params": {
        "url": "file://绝对路径,指向上一步打包的文件.zip"
      }
    },
    "dataAccess": {
      "package": "@pipcook/plugins-pascalvoc-data-access"
    },
    "dataProcess": {
      "package": "@pipcook/plugins-tfjs-image-classification-process",
      "params": {
        "resize": [224, 224]
      }
    },
    "modelDefine": {
      "package": "@pipcook/plugins-tfjs-mobilenet-model-define",
      "params": {}
    },
    "modelTrain": {
      "package": "@pipcook/plugins-image-classification-tfjs-model-train",
      "params": {
        "batchSize": 64,
        "epochs": 12
      }
    },
    "modelEvaluate": {
      "package": "@pipcook/plugins-image-classification-tfjs-model-evaluate"
    }
  }
}
复制代码

使用 Pipcook 配套的 Cli 工具开始训练:

$ pipcook run 上面写的配置项.json
复制代码

看到出现 Epochs 和 Iteration 字样说明训练成功开始了。

...
ℹ [job] running modelTrain start
ℹ start loading plugin @pipcook/plugins-image-classification-tfjs-model-train
ℹ @pipcook/plugins-image-classification-tfjs-model-train plugin is loaded
ℹ Epoch 0/12 start
ℹ Iteration 0/303 result --- loss: 5.969481468200684 accuracy: 0
ℹ Iteration 30/303 result --- loss: 5.65574312210083 accuracy: 0.015625
ℹ Iteration 60/303 result --- loss: 5.293442726135254 accuracy: 0.0625
ℹ Iteration 90/303 result --- loss: 4.970404624938965 accuracy: 0.03125
...
复制代码

两万多张样本以上面的参数在我的 Mac 上训练大约需要两个小时,期间电脑的 cpu 资源都会被占用,所以要找好空闲的时间训练。如果中途要停下来,用 control + c 是没用的,需要先用 pipcook job list 查看任务列表,再用 pipcook job stop <jobId> 来停止训练。

训练的时长与:样本的数据量、epochs 和 batchSize 有关。

/* =============== 两个小时后... =============== */

训练完成,能看到最终的损失率(越低越好)和准确率(越高越好):

...
ℹ [job] running modelEvaluate start
ℹ start loading plugin @pipcook/plugins-image-classification-tfjs-model-evaluate
ℹ @pipcook/plugins-image-classification-tfjs-model-evaluate plugin is loaded
ℹ Evaluate Result: loss: 0.05339580587460659 accuracy: 0.9850694444444444
...
复制代码

如果损失率大于 0.2,准确率低于 0.8,那训练的效果就不太好了,需要调整参数或样本,然后重新训练。

同时 pipcook 会在配置项 json 同目录下创建一个 output 文件夹,里面包含了我们需要的模型:

output
  |-- logs    # 训练日志文件夹
  |-- model   # 模型文件夹,里面两个文件就是最终需要的产物
      |-- weights.bin
      |-- model.json
  |-- metadata.json    # 元信息
  |-- package.json     # 项目信息
  |-- index.js         # 默认入口文件
  |-- boapkg.js        # 辅助文件
复制代码

模型使用

因为用的 Pipcook 插件底层调用 Tensorflow.js 进行训练,所以模型可以直接在前端页面运行。

我们先把生成的 model.jsonweights.bin 放在同一目录下存好。然后找到 metadata.json 中的 output.dataset 字段,是个 Json 字符串,反序列化后找到的 labelArray 属性的值并且存下来:

// 目前这个顺序是随机生成的,和样本生成时的顺序不一样,不要混淆了
const labelArray = ["col-before","h1","solidDown","add-test",...];
复制代码

准备就绪,只要再写一些 Tensorflow.js 代码就可以进行识别了。

import * as tf from '@tensorflow/tfjs';

const modelUrl = 'model.json 的访问地址';
// 加载模型
model = await tf.loadLayersModel(modelUrl);

// 对输入图像裁剪
const { x, y, width: w, height: h } = getCutPosition(imgW, imgH, offscreenCtx.getImageData(0, 0, imgW, imgH).data, 'white');
ctx.drawImage(offscreenCanvas, x, y, w, h, 0, 0, cutSize, cutSize);

// 图像转化为 tensor
const imgTensor = tf.image
  .resizeBilinear(tf.browser.fromPixels(canvas), [224, 224])
  .reshape([1, 224, 224, 3]);
  
// 模型识别
const pred = model.predict(imgTensor).arraySync()[0];

// 找出相似度最高的 5 项
const result = pred.map((score, i) => ({ score, label: labelArray[i] }))
  .sort((a, b) => b.score - a.score)
  .slice(0, 5);
复制代码

大功告成

Mar-21-2021 23-20-30.gif

现在可以开始体验图标识别的能力,享受机器学习带来的便利了。这是一个纯前端工具,无需额外后端服务,可以在静态网站上部署,非常适合在组件库网站中查找图标的场景。团队有自己的图标库也完全没问题,只要按照步骤走,就能训练出专属的模型。

完整代码见:github.com/maplor/icon…

总结

从开始写代码到模型能用花了一个周末加两个晚上,而搭建环境和训练模型的时间占了很大比例。Pipcook 虽然使用简单,省去了很多工作,但入门也有不少坑:文档稀少,插件的参数只有看源码才明白,运行过程有一些潜规则需要不断试错。希望 Pipcook 的文档能及时更新和维护。

如果有什么疑问可以在评论指出,欢迎大家体验交流~

常见问题

  • 图标库如果有 新增/修改 图标怎么办?答:需要重新训练模型。

参考资料

斯坦福《机器学习》课程

《Tensorflow.js 海量图标,毫秒级识别!》

Tensorflow.js 官网

Pipcook 官网

一文看懂机器学习

一文看懂卷积神经网络 CNN

加入我们

我们是阿里云的 TXD(体验技术)团队,诚招前端和设计师,22届的实习生校招也在火热进行中,感兴趣的同学可以联系我了解更多信息:zhaoye.zzy@alibaba-inc.com

mmexport1616514992052.jpg

文章分类
前端