前端如何用AI实现证件照在线换底色

点击上方 前端瓶子君，关注公众号

回复算法，加入前端编程面试算法每日一题群

• 预览地址：https://zhangyuyan.cn/example/replace-background

• 源码：https://github.com/RyanProMax/interest/blob/main/src/components/ReplaceBackground

最近学机器学习刚好学到 k-means 聚类算法，心血来潮便写个 demo 作为实践。

原理很简单：

1. 读取图片数据，将其转换为矩阵数据集，单位是像素，特征值为：r, g, b, a。

2. 通过 k-means 算法对数据集进行分类。

3. 拿到背景色对应的分类，将该分类下所有像素颜色转换成目标色，最终输出图像即可。

1. k-means 是什么？

维基百科：k-平均演算法[3]

【机器学习】K-means（非常详细）[4]

这是一种基于欧式距离的聚类算法。

假设二维平面上有一个点集，我们现在想把这些点分为 3 类（k=3）。

image.png

你可能会说，这不是很明显嘛？将三个方块划区就好啦。

可是机器不同于人的思维，我们必须先给他指定划分区域的规则，才有可能使其实现自动分类。

那么区域有什么规则特性呢？

假设每个区域都有一个区域中心（聚类中心）的话，那么该区域内所有的点，与其对应的区域中心距离最近。

换言之，我们只要找到三个最好的区域中心，那么便得到属于该区域的点集。整个聚类问题便转化为：计算区域中心的最优解问题。

那么区域中心什么情况下才算是最优呢？

当区域中心近似于该区域点集的平均中心时。

算法步骤

1. 在随机位置上初始化 k 个点作为【聚类中心】；

2. 针对数据集中每个样本，计算它到 k 个聚类中心的距离，并将其分到距离最小的聚类中心所对应的类中；

3. 针对每个类别的数据集，重新计算它的聚类中心，即属于该类的所有样本的质心；

4. 重复 2、3 步骤，直到聚类中心达到最小误差。

我们首先随机生成 3 个区域中心，并找到与之对应距离最短的点集（区域）。可以看到随机分类结果很差。

image.png

然后我们先计算以上区域的平均中心，并作为新的区域中心去重新计算各区域点集。

此时结果如下，可以看到已经很接近我们想要的结果了。

image.png

最后循环代入，便得到了我们最终的分类结果。

k-means.gif

• 在线演示地址：ryanpromax.github.io/MachineLear…[5]

• 在线演示源码：github.com/RyanProMax/…[6]

2. 在线换底色程序

2.1 读取图片数据，转换为矩阵数据集

// 获得图片的宽高以及 img 标签
const getImageSize = img => {
  return new Promise((resolve, reject) => {
    const image = new Image();
    image.src = img;
    image.onload = e => {
      if (e.path && e.path.length) {
        const { width, height } = e.path[0];
        resolve([width, height, image]);
      } else {
        reject();
      }
    };
    image.onerror = reject;
  });
};

const transition = imageData => {
  const ret = [];
  for (let i = 0; i < imageData.length; i += 4) {
    ret.push([imageData[i], imageData[i + 1], imageData[i + 2], imageData[i + 3]]);
  }
  return ret;
};

// 通过 canvas 拿到 imageData，并将其转化为 rgba 像素矩阵
export const getImageData = async img => {
  const [width, height, image] = await getImageSize(img);
  const canvas = document.createElement('canvas');
  [canvas.width, canvas.height] = [width, height];
  const ctx = canvas.getContext('2d');
  ctx.drawImage(image, 0, 0);
  return [transition(ctx.getImageData(0, 0, width, height).data), width, height];
};
复制代码

返回结果第一个是二维矩阵（数组），row 是每个像素，column 为像素对应的 r, g, b, a 值，即特征向量。

后面两个分别是图片的宽高。

image.png

2.2 k-means 分类

这里懒得自己写（狗头.jpg），就用了 node-kmeans 库，将像素集分为四类（其实也可以分为两类，但个人觉得四类会精确点，即可能减少其他颜色的干扰。但相对应的，计算时间也会增加，按需设置吧）。

const kmeans = require('node-kmeans');

export const useKmeans = vectors => {
  return new Promise((resolve, reject) => {
    kmeans.clusterize(vectors, { k: 4 }, (err, res) => {
      if (err) reject(err);
      resolve(res);
    });
  });
};
复制代码

返回结果形式如下：

image.png

• centroid 是聚类中心的特征值数组。

• cluster 是类对应的点集。

• clusterInd 是 cluster 对应的下标。

这里直接取第 0 个像素所属的区域为背景色区域。

const result = await useKmeans(imageData);
const target = result.find(x => x.clusterInd.includes(0));
复制代码

（别问如果证件照有边框怎么办，问就是懒，狗头.jpg）

2.3 转换背景色对应像素的颜色

这里就不再赘述了，就是把像素的 rgba 值改为目标颜色值即可，然后再将 rgba 像素矩阵转换回图片输出即可。

个人是通过 canvas 转换的，性能不一定最优，不喜勿喷hh。

export const toBase64 = (imageData, w, h) => {
  imageData = imageData.flat(1);
  const canvas = document.createElement('canvas');
  [canvas.width, canvas.height] = [w, h];
  const ctx = canvas.getContext('2d');
  const _imageData = ctx.createImageData(w, h);
  if (_imageData.data.set) {
    _imageData.data.set(imageData);
  } else {
    // IE9
    imageData.forEach(function (val, i) {
      _imageData.data[i] = val;
    });
  }
  ctx.putImageData(_imageData, 0, 0);
  return canvas.toDataURL();
};
复制代码

至此，咱们就转换完成了。

别问我如果证件照人像处内有干扰色（跟背景色相近，噪点）怎么办，这个留给大家自由发挥（个人的想法是通过连通域算法来解决，但主要是懒得写）。

END~

来自：RyanProMax

https://juejin.cn/post/7051509977269141535

最后

欢迎关注【前端瓶子君】✿✿ヽ(°▽°)ノ✿

回复「算法」，加入前端编程源码算法群，每日一道面试题（工作日），第二天瓶子君都会很认真的解答哟！

回复「交流」，吹吹水、聊聊技术、吐吐槽！

回复「阅读」，每日刷刷高质量好文！

如果这篇文章对你有帮助，「在看」是最大的支持

 》》面试官也在看的算法资料《《

“在看和转发”就是最大的支持