extractColorFromImage.ts

// From internal database of SCG Apply(2024)
// by lhwdev 2023.

// Modified code from https://github.com/luukdv/color.js (modified a lot)
// License: MIT License
// (https://github.com/luukdv/color.js/blob/eba683b1c9197d0fd8cf460cf251f297a94e0bf3/license.md)


import { LectureThumbnail } from "@/components/pages/dashboard/Lecture";
import { Color } from "@/theme/color";

type Args = {
  amount: number;
  format: string;
  group: number;
  sample: number;
  maximumPoints: number;

  alphaThreshold: number;

  debugName?: string;
};

export interface Data {
  width: number;
  height: number;
  pixels: Uint8Array | Uint8ClampedArray;
}
type Item = HTMLImageElement | Data;
type Handler<T> = (data: Data, args: Args) => T;

const calculateSample = (pixels: number, sample: number, maximumPoints: number) => {
  const iterations = Math.floor(pixels / sample);
  if (iterations > maximumPoints * maximumPoints) {
    return Math.ceil(pixels / (maximumPoints * maximumPoints));
  }
  return sample;
};

const getArgs = ({
  amount = 3,
  format = "array",
  group = 32,
  sample = 10,
  maximumPoints = 54,
  alphaThreshold = 250,
  ...args
}: Partial<Args> = {}): Args => ({
  amount,
  format,
  group,
  sample,
  maximumPoints,
  alphaThreshold,
  ...args,
});

const group = (number: number, grouping: number): number => {
  const grouped = Math.round(number / grouping) * grouping;

  return Math.min(grouped, 255);
};

const getImageData = (image: HTMLImageElement): Data => {
  const canvas = document.createElement("canvas");
  const context = canvas.getContext("2d")!;
  canvas.height = image.height;
  canvas.width = image.width;
  context.drawImage(image, 0, 0);

  return {
    width: canvas.width,
    height: canvas.height,
    pixels: context.getImageData(0, 0, image.width, image.height).data,
  };
};

const getAverage: Handler<Color[]> = (data, args) => {
  const { pixels } = data;
  const gap = 4 * args.sample;
  let amount = 0;
  const rgb = { r: 0, g: 0, b: 0 };

  for (let i = 0; i < pixels.length; i += gap) {
    if (pixels[i + 3] < args.alphaThreshold) continue;
    rgb.r += pixels[i];
    rgb.g += pixels[i + 1];
    rgb.b += pixels[i + 2];
    amount++;
  }

  return [
    Color.fromRgbaIntComponents(
      Math.round(rgb.r / amount),
      Math.round(rgb.g / amount),
      Math.round(rgb.b / amount)
    ),
  ];
};

const getProminent: Handler<Color[]> = ({ pixels }, args) => {
  const gap = 4 * args.sample;
  const colors = new Map<number, number>();

  for (let i = 0; i < pixels.length; i += gap) {
    if (pixels[i + 3] < args.alphaThreshold) continue;
    const rgb =
      ((group(pixels[i], args.group) << 16) |
        (group(pixels[i + 1], args.group) << 8) |
        group(pixels[i + 1], args.group)) >>>
      0;

    const previous = colors.get(rgb) ?? 0;
    colors.set(rgb, previous + 1);
  }

  return Array.from(colors.entries())
    .sort(([_keyA, valA], [_keyB, valB]) => (valA > valB ? -1 : 1))
    .slice(0, args.amount)
    .map(([rgb]) => Color.fromRgbInt(rgb));
};

const debugMeasureTiming = false;
const debugMeasureTimingDetail = debugMeasureTiming;
const getLectureThumbnail: Handler<Omit<LectureThumbnail, "imageUrl">> = (
  { width, height, pixels },
  args
) => {
  if (debugMeasureTiming) {
    console.time("getLectureThumbnail");
  }
  if (debugMeasureTimingDetail) console.time(":phase1");
  const sample = calculateSample(width * height, args.sample, args.maximumPoints);

  const paletteCount = 3;
  const gap = 4 * sample;

  // note: rank is here is also used to measure count until some point in this code
  const colors = new Map<number, [rank: number, r: number, g: number, b: number]>();
  let transparent = 0;
  const amount = pixels.length / gap;
  const sum = [0, 0, 0];
  let edgeOpaque = 0;

  for (let i = 0; i < pixels.length; i += gap) {
    const a = pixels[i + 3];
    if (a < args.alphaThreshold) {
      transparent++;
      continue;
    }
    const r = pixels[i];
    const g = pixels[i + 1];
    const b = pixels[i + 2];

    const x = i % width;
    const y = Math.floor(i / width);
    if (x === 0 || x === width - 1 || y === 0 || y === height - 1) {
      // a >= args.alphaThreshold already
      edgeOpaque++;
    }

    // Find out color with maximum appearance
    const rgb =
      ((group(r, args.group) << 16) | (group(g, args.group) << 8) | group(b, args.group)) >>> 0;
    const [prevRank, prevR, prevG, prevB] = colors.get(rgb) ?? [0, 0, 0, 0];
    // NOTE: they might overflow; i don't care
    colors.set(rgb, [prevRank + 1, prevR + r, prevG + g, prevB + b]);

    // Average color
    sum[0] += r;
    sum[1] += g;
    sum[2] += b;
  }
  if (debugMeasureTimingDetail) console.timeEnd(":phase1");
  if (debugMeasureTimingDetail) console.log(`- ${pixels.length / gap} ${colors.size}`);
  if (debugMeasureTimingDetail) console.time(":phase2");

  const results = Array.from(colors.entries());
  let result: [color: number, rank: number, isVibrant: boolean][] = Array(results.length);

  result.sort((a, b) => a[1] - b[1]);
  result = result.slice(0, 8); // if there are too much colors, cut them out early

  let vibrantCount = 0;
  const globalAverage = sum.map((n) => n / amount);

  for (let i = 0; i < results.length; i++) {
    const [_, [rank, averageR, averageG, averageB]] = results[i];
    let r = averageR / rank;
    let g = averageG / rank;
    let b = averageB / rank;

    // Apply global average
    const lerp = (a: number, b: number) => Math.round(a + (b - a) * 0.1);
    r = lerp(r, globalAverage[0]);
    g = lerp(g, globalAverage[1]);
    b = lerp(b, globalAverage[2]);

    const average = ((r << 16) | (g << 8) | b) >>> 0;

    // just a standard deviation of r, g, b
    const medium = (r + g + b) / 3;
    const devR = r - medium;
    const devG = g - medium;
    const devB = b - medium;
    const vibrance = Math.sqrt(devR * devR + devG * devG + devB * devB);
    // Prefer more vibrant colors
    // If too monotonic, remove from palette
    const isVibrant = vibrance >= 5;
    result[i] = [average, rank * Math.max(1, 1.5 - 1 / (vibrance * 0.2 + 2)), isVibrant];
    if (isVibrant) vibrantCount++;
  }

  let mapped: [color: Color, rank: number, isVibrant: boolean][] = result.map(
    ([color, rank, isVibrant]) => [Color.fromRgbInt(color), rank, isVibrant]
  );
  const swap = (a: number, b: number) => {
    const temp = mapped[a];
    mapped[a] = mapped[b];
    mapped[b] = temp;
  };
  const findMax = (previous: Color | null, fromIndex: number) => {
    const colorDistance = (a: Color, b: Color) => {
      return a.distanceTo(b);
    };
    let max = fromIndex;
    let maxRank = 0;
    for (let i = fromIndex; i < mapped.length; i++) {
      const [color, rank, _isVibrant] = mapped[i];
      let realRank = rank;
      if (previous) {
        realRank *= Math.pow(colorDistance(previous, color), 0.4) + 1;
      }
      if (maxRank < realRank) {
        max = i;
        maxRank = realRank;
      }
    }
    return max;
  };
  for (let i = 0; i < mapped.length - 1; i++) {
    const maxIndex = findMax(i === 0 ? null : mapped[i - 1][0], i);
    swap(i, maxIndex);
  }

  if (vibrantCount < paletteCount) {
    mapped = mapped.slice(0, paletteCount);
  } else {
    mapped = mapped.filter(([, , isVibrant]) => isVibrant);
  }

  const edgeCount = 2 * (width + height) - 4;
  const type = edgeOpaque / edgeCount < 0.1 && transparent / amount > 0.1 ? "logo" : "backdrop";

  const colorAt = (index: number) => (index >= mapped.length ? mapped[0] : mapped[index])[0];

  if (debugMeasureTimingDetail) console.timeEnd(":phase2");
  if (debugMeasureTimingDetail) console.time(":phase3");
  const finalResult: Omit<LectureThumbnail, "imageUrl"> = {
    type,
    keyColors: {
      primary:
        type === "logo"
          ? colorAt(0).withHct((hct) => {
              hct.tone =
                hct.tone < 50 ? Math.max(100 - hct.tone, 90) : Math.min(100 - hct.tone, 10);
              hct.chroma *= 0.5;
            })
          : colorAt(0).withHct((hct) => {
              const minimum = Math.abs(hct.tone - 50) / 5 + 10;
              const newChroma = hct.chroma * 0.8;
              hct.chroma = hct.chroma <= minimum ? hct.chroma : Math.max(minimum, newChroma);
            }),
      secondary: colorAt(1),
      tertiary: colorAt(2),
    },
  };
  if (debugMeasureTimingDetail) console.timeEnd(":phase3");
  if (debugMeasureTiming) {
    console.timeEnd("getLectureThumbnail");
    console.log(`- ${args.debugName} sample=${sample} size=${width}*${height}`);
    // console.log({
    //   type: finalResult.type,
    //   keyColors: Object.fromEntries(
    //     Object.entries(finalResult.keyColors).map(([k, v]) => [k, v.value])
    //   ),
    // });
  }
  return finalResult;
};

const process = <T>(handler: Handler<T>, item: Item, args?: Partial<Args>): T =>
  handler(
    typeof HTMLImageElement === "undefined"
      ? (item as Data)
      : item instanceof HTMLImageElement
      ? getImageData(item)
      : item,
    getArgs(args)
  );

export const average = (item: Item, args?: Partial<Args>) => process(getAverage, item, args);
export const prominent = (item: Item, args?: Partial<Args>) => process(getProminent, item, args);
export const lectureThumbnail = (item: Item, args?: Partial<Args>) =>
  process(getLectureThumbnail, item, args);