color quantization

get-image-data.ts

export const getImageData = (image: HTMLImageElement) => {
  const canvas = document.createElement('canvas');
  const ctx = canvas.getContext('2d') as CanvasRenderingContext2D;
  const width = image.width;
  const height = image.height;

  if (width === 0 || height === 0) throw new Error('image dimensions cannot be zero');

  canvas.width = width;
  canvas.height = height;
  ctx.drawImage(image, 0, 0, width, height);
  return ctx.getImageData(0, 0, width, height);
};

median-cut.ts

import { getImageData } from './get-image-data';

export type RGBPixel = number[];
export type RGBData = RGBPixel[];
export type RGBChannels = number[][];

export const getRGBChannels = (rgbData: RGBData): RGBChannels => {
  const rChannel: number[] = [];
  const gChannel: number[] = [];
  const bChannel: number[] = [];
  rgbData.forEach((rgb) => {
    const [r, g, b] = rgb;
    rChannel.push(r);
    gChannel.push(g);
    bChannel.push(b);
  });

  return [rChannel, gChannel, bChannel] as RGBChannels;
};

export const getRGBData = (imageData: Uint8ClampedArray) => {
  const rgbData: number[][] = [];
  // srgb has 4 channels
  for (let i = 0; i < imageData.length; i += 4) {
    rgbData.push(Array.prototype.slice.call(imageData, i, i + 3)); // rgb only
  }
  return rgbData;
};

export const splitColors = (rgbData: RGBData, depth: number): RGBData[] => {
  if (!depth) return [rgbData];

  const [rChannel, gChannel, bChannel] = getRGBChannels(rgbData);
  const channelRanges = [
    Math.max(...rChannel) - Math.min(...rChannel),
    Math.max(...gChannel) - Math.min(...gChannel),
    Math.max(...bChannel) - Math.min(...bChannel),
  ];

  const channelLength = rChannel.length;
  const maxRange = Math.max(...channelRanges);
  const maxRangeChannelIndex = channelRanges.indexOf(maxRange);
  const channelMedian = rgbData.reduce((acc, x) => acc + x[maxRangeChannelIndex], 0) / channelLength;
  const bottom: RGBData = [];
  const top: RGBData = [];

  rgbData.forEach((d) => {
    if (d[maxRangeChannelIndex] < channelMedian) {
      bottom.push(d);
    } else {
      top.push(d);
    }
  });

  return [...splitColors(bottom, depth - 1), ...splitColors(top, depth - 1)];
};

export const blockMeans = (rgbData: RGBData) => {
  const length = rgbData.length;
  return rgbData.reduce(
    (acc, pixel) => {
      pixel.forEach((c, i) => (acc[i] += c / length));
      return acc;
    },
    [0, 0, 0]
  );
};

export const medianCut = (rgbData: RGBPixel[], buckets: number) => {
  const chunks = splitColors(rgbData, Math.log2(buckets));
  return chunks.map(blockMeans);
};

export const extractColor = (img: HTMLImageElement) => {
  const imageData = getImageData(img);
  const rgbData = getRGBData(imageData.data);
  const result = medianCut(rgbData, 4);
  return result;
};