Batch recolor an image to match a series of provided templates.

batchRecolor.js

const { loadImage, createCanvas } = require("@napi-rs/canvas");
const { readdirSync, writeFileSync, existsSync, mkdirSync } = require("fs");
const { join } = require("path");

/** @typedef {string | URL | Buffer | ArrayBufferLike | Uint8Array | import("@napi-rs/canvas").Image | import("stream").Readable} ImageSource */

/** @typedef {[number, number, number, number]} ColorTuple - [r, g, b, a] array of colors */

/**
 * @typedef ReplacementData
 * @property {ColorTuple} original - Original color on image
 * @property {ColorTuple} replacement - Color to replace with
 */

/** Quick helper method to compare two fixed-sized ordered arrays */
const tupleEqual = (a, b) => JSON.stringify(a) === JSON.stringify(b);

/**
 * Nicer formatted ImageData for getting data
 * @param {import("@napi-rs/canvas").SKRSContext2D} ctx - Where to get data from
 * @returns {ColorTuple[]} Array of [r, g, b, a] tuples
 */
const getImageDataAsTuples = (ctx) =>
	ctx.getImageData(0, 0, ctx.canvas.width, ctx.canvas.height).data.reduce((acc, cur, i) => {
		if (i % 4 === 0) acc.push([]);
		acc.at(-1).push(cur);
		return acc;
	}, []);

/**
 * Helper method to load directly as a Canvas rather than as an Image
 * @param {ImageSource} source - Any loadable image
 * @returns {Promise<import("@napi-rs/canvas").SKRSContext2D>} Loaded canvas context
 */
async function loadToCanvas(source) {
	const image = await loadImage(source);
	const canvas = createCanvas(image.width, image.height);
	const ctx = canvas.getContext("2d");
	ctx.imageSmoothingEnabled = false;
	ctx.drawImage(image, 0, 0);
	return ctx;
}

/**
 * Generate a pixel color map of unique images
 * @param {ImageSource} reference - Image to get colors of
 * @param {ImageSource} image - Images to check colors exist in
 * @returns {Promise<Record<number, ColorTuple>>} Generated pixel color map
 */
async function generatePixelMap(reference, image) {
	const ctx = await loadToCanvas(image);
	const refCtx = await loadToCanvas(reference);

	const imageData = getImageDataAsTuples(ctx);
	const refData = getImageDataAsTuples(refCtx);

	console.log(`Generating pixel map for ${refData.length} possible colors...`);
	return refData.reduce((acc, data, i) => {
		if (
			!data[3] ||
			!imageData.some((d) => tupleEqual(d, data)) ||
			Object.values(acc).some((el) => tupleEqual(el, data))
		)
			return acc;
		return { ...acc, [i]: data };
	}, []);
}

/**
 * Convert a pixel color map into a replacement color map
 * @param {Record<number, ColorTuple>} pixelMap - Pixel map to load
 * @param {ImageSource} source - Image to generate for
 * @returns {Promise<ReplacementData[]>} Replacement color map
 */
async function generateReplacementMap(pixelMap, source) {
	const template = await loadToCanvas(source);
	const imageData = getImageDataAsTuples(template);

	// can't use array as key so object/map can't be used
	return imageData.reduce((acc, data, i) => {
		const match = pixelMap[i];
		if (!match) return acc;
		return [...acc, { original: match, replacement: data }];
	}, []);
}

/**
 * Apply a generated replacement map to a given image
 * @param {ReplacementData[]} replacementMap - Replacement map to apply
 * @param {import("@napi-rs/canvas").SKRSContext2D} source - Image to edit
 * @returns {Promise<import("@napi-rs/canvas").SKRSContext2D>} Edited image
 */
async function applyReplacementMap(replacementMap, source) {
	const image = await loadToCanvas(source);

	// will be putting back so we need original format
	const imageData = image.getImageData(0, 0, image.canvas.width, image.canvas.height);
	for (let i = 0; i < imageData.data.length; i += 4) {
		const matched = replacementMap.find((el) =>
			tupleEqual(el.original, Array.from(imageData.data.slice(i, i + 4))),
		);
		if (matched) {
			imageData.data[i] = matched.replacement[0];
			imageData.data[i + 1] = matched.replacement[1];
			imageData.data[i + 2] = matched.replacement[2];
			imageData.data[i + 3] = matched.replacement[3];
		}
	}

	image.putImageData(imageData, 0, 0);
	return image;
}

/**
 * Batch recolor an image to match a series of provided templates
 * @author Evorp
 * @param {{ reference: string, image: string, templateDir: string, outputDir: string }} params - Everything to use
 * @returns {Promise<void>} Writes the recolored images to the provided outputDir
 */
async function batchRecolor({ reference, image, templateDir, outputDir }) {
	const pixelMap = await generatePixelMap(reference, image);
	console.log(`Successfully reduced pixel map to ${Object.keys(pixelMap).length} colors!`);

	if (!existsSync(outputDir)) mkdirSync(outputDir, { recursive: true });
	for (const file of readdirSync(templateDir).filter((i) => i.endsWith(".png"))) {
		console.log(`Recoloring ${image} to ./${file}...`);

		const replacementMap = await generateReplacementMap(pixelMap, join(templateDir, file));
		const recoloredImage = await applyReplacementMap(replacementMap, image);

		writeFileSync(join(outputDir, file), recoloredImage.canvas.toBuffer("image/png"));
	}
}

// sample usage
batchRecolor({
	reference: "./ref.png",
	image: "./red.png",
	templateDir: "./templates/",
	outputDir: "./output/",
});