Extract Prominent Colors from an Image Using Machine Learning

Ok, I did this back in 2014 in java but since I reimplemented the app and the algorithm in javascript this time. I thought of writing about how I found a nice and simple algorithm to extract prominent colors out of an image.

Problem:

Given an image we want to extract 6 prominent colors from it.

Definitions:

LAB color space: Lab color is designed to approximate human vision. Euclidean distance in LAB color space approximates perceivable distance of human vision.

K-means: It is unsupervised clustering algorithm that groups the unlabeled dataset into k different clusters.

HSB(Hue Saturation Brightness) color space

Algorithm:

1. Resize the image to a manageable size (100 x 100)

2. Convert individual pixels of colors to LAB color space

3. Calculate 6 * 4 = 24 k means clusters from pixels in LAB color space

4. Convert these 24 centroids to HSB color space and sort them by hue.

5. Divide them into groups of 4 and for each group pick 2 most saturated colors

6. From 2 colors pick the most bright color

Here is the full code.

import kmeans from "ml-kmeans";
import Jimp from "jimp";
import Color from "pigment/full";
export default class ColorPicker {
  static getProminentColors(image) {
    /*
    Jimp.RESIZE_NEAREST_NEIGHBOR;
    Jimp.RESIZE_BILINEAR;
    Jimp.RESIZE_BICUBIC;
    Jimp.RESIZE_HERMITE;
    Jimp.RESIZE_BEZIER;
    These does not work with first params.
    */
    image.resize(Jimp.AUTO, 100);
    let data = ColorPicker._prepareDataForKmeans(image);
    let time = Date.now();
    let ans = kmeans(data, 24, { initialization: "random", maxIterations: 20 });
    ans.centroids = ans.centroids.sort((c1, c2) => c2.size - c1.size);
    let kmeansColors = ans.centroids.map(centroid => {
      return new Color(this._labToHex(centroid.centroid));
    });
    return this._getFinalColors(kmeansColors).map(c => {
      return { color: c.tohex() };
    });
  }
  // original implementation in java: https://github.com/kamalkishor1991/croma/blob/master/src/main/java/org/numixproject/colorextractor/image/KMeansColorPicker.java
  static _getFinalColors(kmeansColors) {
    kmeansColors.sort(
      (c1, c2) => this._toArray(c1.tohsv())[0] < this._toArray(c2.tohsv())[0]
    );
    let filteredColors = [];
    for (let i = 0; i < kmeansColors.length; i += 4) {
      let colorList = [];
      for (let j = 0; j < 4; j++) {
        colorList.push(kmeansColors[i + j]);
      }
      colorList.sort(
        (c1, c2) => this._toArray(c1.tohsv())[1] < this._toArray(c2.tohsv())[1]
      );
      filteredColors.push(colorList[colorList.length - 1]);
      filteredColors.push(colorList[colorList.length - 2]);
    }
    let finalColors = [];
    for (let i = 0; i < filteredColors.length; i += 2) {
      if (
        this._toArray(filteredColors[i].tohsv())[2] >
        this._toArray(filteredColors[i + 1].tohsv())[2]
      ) {
        finalColors.push(filteredColors[i]);
      } else {
        finalColors.push(filteredColors[i + 1]);
      }
    }
    return finalColors;
  }

  static _labToHex(lab) {
    let color = new Color(
      "lab(" + lab[0] + ", " + lab[1] + ", " + lab[2] + ")"
    );
    return color.tohex();
  }

  static _prepareDataForKmeans(image) {
    let data = [];

    for (let i = 0; i < image.bitmap.width; i++) {
      for (let j = 0; j < image.bitmap.height; j++) {
        let intColor = image.getPixelColor(i, j);
        let hex = this._toHexColor(intColor);
        let color = new Color(hex);
        let xyz = color.tolab();
        // format: "xyz(19.78527130484015, 8.600439447528947, 95.19796416837329)" to double array of xyz
        xyz = xyz
          .substr(4, xyz.length - 5)
          .split(", ")
          .map(v => parseFloat(v));
        data.push(xyz);
      }
    }

    return data;
  }

  static _toHexColor(intColor) {
    let rgba = Jimp.intToRGBA(intColor); // TODO: Need to optimize this once everything else starts working.
    let color = new Color(
      "rgb(" + rgba.r + ", " + rgba.g + ", " + rgba.b + ")"
    );
    return color.tohex();
  }

  static _toArray(color) {
    let index = color.indexOf("(");
    color = color.substr(index + 1, color.length - index);
    return color.split(", ").map(c => parseFloat(c));
  }
}

Javascript implementation

Java implementation

Try it out here - https://croma.app or on the Google Play Store

Process and Learnings:

It turns out that this gives us really good results and off-course I tried a lot of other methods before reaching to this algorithm. Here are some of the key properties which make this algorithm so good.

1. Converting to LAB color space - This divides the color space of the image into an evenly distributed space based on human perception of colors. I tried k-means with other color spaces like RGB but they did not yield good results as their Euclidean distance does not make much sense to human eye.

2. Resizing the image to avoid performance issues with k-means as it is a n^2 algorithm so reducing the number of points to execute it fast on mobile device is critical.

3. Getting 24 points and further sorting them based on HSB.

I had fun developing this back is 2014 and learned a lot about color spaces and developing simple algorithm to solve problems.

My biggest learning while developing this was

Elegant solutions require a deep understanding of the problem space.

Thanks for reading and happy coding. The app is open source so please consider contributing and making it better.