refactor: improve performance and accuracy for images with few colors

Author: ksubileau

src/pipeline.ts

@@ -188,9 +188,39 @@ export function extractPalette(
         });
     }
 
-    // OKLCH quantization path
     let quantized: Array<{ color: [number, number, number]; population: number }>;
-    if (opts.colorSpace === 'oklch') {
+
+    // Short-circuit check: collect up to opts.colorCount + 1 unique colors
+    const seenColors = new Set<string>();
+    const uniqueColors: Array<[number, number, number]> = [];
+
+    for (const color of pixelArray) {
+        const key = color.join(',');
+        if (!seenColors.has(key)) {
+            seenColors.add(key);
+            uniqueColors.push(color);
+
+            // The image contains more distinct colors than requested,
+            // so we can stop searching for unique colors as we know we will have to use the quantizer.
+            if (uniqueColors.length > opts.colorCount) break;
+        }
+    }
+
+    // If unique colors <= maxColors, return them directly with population counts
+    if (uniqueColors.length <= opts.colorCount) {
+        // Count populations for unique colors
+        const countMap = new Map<string, number>();
+        for (const color of pixelArray) {
+            const key = color.join(',');
+            countMap.set(key, (countMap.get(key) || 0) + 1);
+        }
+        quantized = uniqueColors.map((color) => ({
+            color,
+            population: countMap.get(color.join(','))!,
+        }));
+    }
+    // OKLCH quantization path
+    else if (opts.colorSpace === 'oklch') {
         const scaled = pixelsRgbToOklchScaled(pixelArray);
         quantized = paletteOklchScaledToRgb(
             quantizer.quantize(scaled, opts.colorCount),

src/quantizers/mmcq.ts

@@ -346,29 +346,6 @@ function quantize(
 ): Array<{ color: [number, number, number]; population: number }> {
     if (!pixels.length || maxColors < 2 || maxColors > 256) return [];
 
-    // Short-circuit: if unique colors <= maxColors, return them directly
-    const seenColors = new Set<string>();
-    const uniqueColors: Array<[number, number, number]> = [];
-    for (const color of pixels) {
-        const key = color.join(',');
-        if (!seenColors.has(key)) {
-            seenColors.add(key);
-            uniqueColors.push(color);
-        }
-    }
-    if (uniqueColors.length <= maxColors) {
-        // Count populations for unique colors
-        const countMap = new Map<string, number>();
-        for (const color of pixels) {
-            const key = color.join(',');
-            countMap.set(key, (countMap.get(key) || 0) + 1);
-        }
-        return uniqueColors.map((color) => ({
-            color,
-            population: countMap.get(color.join(','))!,
-        }));
-    }
-
     const histo = getHisto(pixels);
 
     // Get the initial vbox from the pixels

test/node-test.js

@@ -60,22 +60,22 @@ describe('getColor()', function() {
         expect(isValidRGB(color)).to.be.true;
     });
 
-    it('returns near-black for black.png', async function() {
+    it('returns perfect black for black.png', async function() {
         const color = await getColor(imgPath('black.png'));
         expect(isColorObject(color)).to.be.true;
-        expect(isCloseTo(color, [0, 0, 0])).to.be.true;
+        expect(color.array()).to.deep.equal([0, 0, 0]);
     });
 
-    it('returns near-red for red.png', async function() {
+    it('returns perfect red for red.png', async function() {
         const color = await getColor(imgPath('red.png'));
         expect(isColorObject(color)).to.be.true;
-        expect(isCloseTo(color, [255, 0, 0])).to.be.true;
+        expect(color.array()).to.deep.equal([255, 0, 0]);
     });
 
-    it('returns valid Color for white.png', async function() {
+    it('returns perfect white for white.png', async function() {
         const color = await getColor(imgPath('white.png'));
         expect(isColorObject(color)).to.be.true;
-        expect(isCloseTo(color, [255, 255, 255])).to.be.true;
+        expect(color.array()).to.deep.equal([255, 255, 255]);
     });
 
     it('returns valid Color for transparent.png', async function() {