基于deno和python3研究提取图片主题色的算法

图像主色提取算法

我们在网易云上听歌, 略加设置就能在能看到这样的效果:

网易云是怎么提取出专辑封面主要颜色的呢
首先, 我们需要思考如何表示一张图片. 图片是由一系列像素点组成的, 最简单的表示图片的方法就是用位图, 也即记录下每个像素点的 rgb 来表示
所以我们可以用一个 width * height * 3 的数组来表示一张图片, 其中 width 和 height 分别表示宽高, 3 代表 r,g,b 三个通道

我们以以下4张图片为例进行说明

将他们分别表示为 rgb 像素点, 以 rgb 作为 xyz 坐标, 标注在三维空间中即是这样:

以上是简要说明, 下面将介绍三种常见的图片主题色提取算法的具体实现 ;P

1 kmeans 聚类算法

1. 首先随机确定 n 个初始点, 作为第0阶段的主色色盘(当然这个主色和真实主色相差了很多), 这 n 个点称为 n 个质心

let startCenters: PixelData[] = new Array(mainColorNumber)
  .fill(0)
  .map(() => {
    return new Array(3)
      .fill(0)
      .map(() => ~~(Math.random() * 255)) as PixelData;
  });

2. 进行迭代, 以改善你的猜测
   1. 根据现有质心的位置对原图像的所有像素进行归类, 离哪个质心最近就设置为哪一类
   2. 根据聚类结果, 算出该类别新的质心(该类所有 r、g、b分别求平均值, 得到的点即为新质心)
   3. 根据算出来的质心, 更新startPoints数组
   4. 重复步骤 1, 直至收敛(每个质心迭代后改变的距离均小于minDist)或达到最大迭代次数iterations

计算距离的方法: 将 r,g,b 分别作为空间坐标系的 x,y,z, 计算两点的欧式距离即可

核心代码如下

while (iterations--) {
  /***** 归类 *****/
  center2Cluster = center2Cluster.map(() => []);
  data.forEach((pixel) => {
    // 对每个像素计算距离哪个质心最近
    const closestCenterIndex = startCenters.reduce(
      (prev, curCenter, centerIndex) => {
        const dist = calcDist(curCenter, pixel);
        return dist < prev.dist ? { centerIndex, dist } : prev;
      },
      { centerIndex: -1, dist: Infinity } // 当前距离最近的中心的号码和距离
    ).centerIndex;
    // 将它加入最近的质心的数组中
    center2Cluster[closestCenterIndex].push(pixel);
  });

  /***** 计算新质心 *****/
  const newStartCenters = center2Cluster.map(
    // 对于每一个中心的集合
    (cluster) =>
      cluster
        .reduce(
          // 分别求该集合中各个点 r, g, b 各通道之和
          (total, pixel) => [
            total[0] + pixel[0],
            total[1] + pixel[1],
            total[2] + pixel[2],
          ],
          [0, 0, 0] as PixelData
        )
        .map(
          // 将结果取平均即为新中心
          (totalChannel) => ~~(totalChannel / cluster.length)
        ) as PixelData
  );

  /***** 判断是否收敛 *****/
  let isSettled = true;
  for (let i = 0; i < mainColorNumber; i++) {
    const moveDist = calcDist(
      newStartCenters[i],
      startCenters[i]
    );
    if (moveDist > minMoveDist) {
      isSettled = false;
      break;
    }
  }
  if (isSettled) break;

  /***** 更新 *****/
  startCenters = newStartCenters;
}

结果如下

2 中位切分

1. 将图片每个像素的 r,g,b 值分别作为 x,y,z 坐标, 标在空间坐标系中
2. 用一个最小的立方体框住所有点
3. 将立方体沿一平面切分, 该平面与立方体最长边方向垂直, 使切分得的两部分包含相同数量的像素点
4. 将分得的小立方体递归地按照 3. 的算法切分, 直至分得立方体个数等于所需颜色数即可
5. 将每个立方体中颜色做平均, 即得到最后的主色

其中, 在进行切分时可对待切割立方体做一个排序, 其中单位体积包含像素点越多的立方体越先被切割, 以提高效率

核心代码如下

interface ToCut {
  data: PixelData[]; // 待切分数据
  density: number; // 每单位体积包含的像素点数量
}

/** 计算所给的 data 数组包含的像素点应切分成哪两段 */
function cutIntoTwo(data: PixelData[]): [ToCut, ToCut] {
  // 找到最小的框
  let minNMax = [
    Infinity, // rMin
    -Infinity, // rMax
    Infinity, // gMin
    -Infinity, // gMax
    Infinity, // bMin
    -Infinity, // bMax
  ];
  data.forEach((pixel) => {
    for (let i = 0; i < 3; i++) {
      minNMax[i] = Math.min(minNMax[i], pixel[i]);
      minNMax[i + 1] = Math.max(minNMax[i], pixel[i]);
    }
  });

  // 找到最长边, 以此判断根据 r or g or b 来切分
  let cutBy = -1,
    maxEdge = -Infinity;
  for (let i = 0; i < 3; i++) {
    const curEdge = minNMax[i + 1] - minNMax[i];
    if (curEdge > maxEdge) {
      maxEdge = curEdge;
      cutBy = i;
    }
  }

  const halfNum = ~~(data.length / 2);
  // 按照 toCut 排序, 前一半放一边, 后一半放一边
  const sortedData = [...data].sort((a, b) => a[cutBy] - b[cutBy]);

  const toCutLeft: ToCut = {
      data: sortedData.slice(0, halfNum),
      density: 0,
    },
    toCutRight: ToCut = {
      data: sortedData.slice(halfNum),
      density: 0,
    };
  let vLeft = 1,
    vRight = 1;
  for (let i = 0; i < 3; i++) {
    if (i === cutBy) {
      vLeft *= sortedData[halfNum][cutBy] - minNMax[i]; // 中位数 - min
      vRight *= minNMax[i + 1] - sortedData[halfNum][cutBy]; // max - 中位数
    } else {
      vLeft *= minNMax[i + 1] - minNMax[i];
      vRight *= minNMax[i + 1] - minNMax[i];
    }
  }
  toCutLeft.density = toCutLeft.data.length / (vLeft || 0.0001);
  toCutRight.density = toCutRight.data.length / (vRight || 0.0001);

  return [toCutLeft, toCutRight];
}

结果如下, 三列分别为原图, kmeans 聚类结果, 中位切分结果

在我看来中位切分效果一般…

3 八叉树算法

这是效果最好的算法, 其思想为用一颗八叉树来存储每个像素点的信息, 边存边对相似的颜色进行聚类(即进行剪枝)

0. 将每个点的 RGB 表示为二进制的一行, 堆叠后将每一列的不同编码对应成数字, 共 8 种组合
   RGB 通道逐列黏合之后的值就是其在某一层节点的子节点

   e.g. 如#FF7800，其中 R 通道为0xFF，也就是255，G 为 0x78 也就是120，B 为 0x00 也就是0。
   接下来我们把它们写成二进制逐行放下，那么就是：
   R: 1111 1111
   G: 0111 1000
   B: 0000 0000
   上述颜色的第一位黏合起来是100(2)，转化为十进制就是 4，所以这个颜色在第一层是放在根节点的第5个子节点当中
   第二位是 110(2) 也就是 6，那么它就是根节点的第5个儿子的第7个儿子

1. 建立一棵空八叉树, 设置一个叶子节点个数上限

2. 依次将像素按 0. 的算法插入树中

   1. 若插入后叶子节点数小于上限, 则什么也不做
   2. 若大于上限, 则对最底层的一个非叶子节点的子节点进行合并, 将其转换为叶子节点 rgb 值的平均数, 并清除其子节点

3. 依此类推, 直到最后插入所有的像素, 所得八叉树的叶子节点即为主色调

该算法的核心在于, 具有兄弟关系的子节点的 rgb 每位最多都只相差 1, 即这些颜色非常接近, 所以合并后可以用更少的主色代替这几个像素的颜色

结果如下, 四列分别为原图, kmeans 聚类结果, 中位切分结果 和 八叉树划分结果

可见八叉树划分的结果已经和网易云相当接近了, 入职网易指日可待(x

参考

• 图片主题色提取算法小结
• 中位切割演算法
• 图像主题色提取算法