Finally doing quantization and dithering in a fast and sensible way

Editor/QuantizedPreviewWindow.cs

@@ -127,7 +127,7 @@ public class QuantizedPreviewWindow : EditorWindow
     private void GenerateQuantizedPreview()
     {
 
-        PSXTexture2D psxTex = PSXTexture2D.CreateFromTexture2D(originalTexture, bpp, false);
+        PSXTexture2D psxTex = PSXTexture2D.CreateFromTexture2D(originalTexture, bpp);
 
         quantizedTexture = psxTex.GeneratePreview();
         vramTexture = psxTex.GenerateVramPreview();

Runtime/ImageProcessing.cs

@@ -1,176 +1,171 @@
-using System;
 using System.Collections.Generic;
-using System.Collections.Specialized;
-using System.Diagnostics;
 using System.Linq;
-using System.Threading.Tasks;
-using Codice.CM.Common;
-using DataStructures.ViliWonka.KDTree;
 using UnityEngine;
 
 
-
 namespace PSXSplash.RuntimeCode
 {
-
+    public class TextureQuantizer
-    public class ImageQuantizer
     {
-        private int _maxColors;
+        public struct QuantizedResult
-        private Vector3[,] _pixels;
-        private Vector3[] _centroids;
-        private KDTree kdTree;
-        private int[,] _assignments;
-        private List<Vector3> _uniqueColors;
-
-        public int Width { get; private set; }
-        public int Height { get; private set; }
-
-        public Vector3[] Palette
         {
-            get => _centroids;
+            public int[,] Indices;
+            public List<Vector3> Palette;
         }
 
-        public int[,] Pixels
+        public static QuantizedResult Quantize(Texture2D texture, int maxColors)
         {
-            get => _assignments;
+            int width = texture.width, height = texture.height;
-        }
+            Color[] pixels = texture.GetPixels();
+            int[,] indices = new int[width, height];
+
+            List<Vector3> uniqueColors = pixels.Select(c => new Vector3(c.r, c.g, c.b)).Distinct().ToList();
+            if (uniqueColors.Count <= maxColors) return ConvertToOutput(pixels, width, height);
+
+            List<Vector3> palette = KMeans(uniqueColors, maxColors);
+            KDTree kdTree = new KDTree(palette);
 
 
-        public void Quantize(Texture2D texture2D, int maxColors)
+            // Floyd-Steinberg Dithering
+            for (int y = 0; y < height; y++)
             {
-            Stopwatch stopwatch = new Stopwatch();
+                for (int x = 0; x < width; x++)
-            stopwatch.Start();
-            Color[] pixels = texture2D.GetPixels();
-
-            Width = texture2D.width;
-            Height = texture2D.height;
-
-            _pixels = new Vector3[Width, Height];
-
-            for (int x = 0; x < Width; x++)
                 {
-                for (int y = 0; y < Height; y++)
+                    Vector3 oldColor = new Vector3(pixels[y * width + x].r, pixels[y * width + x].g, pixels[y * width + x].b);
-                {
+                    int nearestIndex = kdTree.FindNearestIndex(oldColor);
-                    Color pixel = pixels[x + y * Width];
+                    indices[x, y] = nearestIndex;
-                    Vector3 pixelAsVector = new Vector3(pixel.r, pixel.g, pixel.b);
+
-                    _pixels[x, y] = pixelAsVector;
+                    Vector3 error = oldColor - palette[nearestIndex];
+                    PropagateError(pixels, width, height, x, y, error);
                 }
             }
 
-            _maxColors = maxColors;
-            _centroids = new Vector3[_maxColors];
-            _uniqueColors = new List<Vector3>();
-
-
-            FillRandomCentroids();
-
-            bool hasChanged;
-            _assignments = new int[Width, Height];
-
-            do
-            {
-                hasChanged = false;
-                for (int x = 0; x < Width; x++)
-                {
-                    for (int y = 0; y < Height; y++)
-                    {
-                        Vector3 color = _pixels[x, y];
-                        int newAssignment = GetNearestCentroid(color);
-                        
-                        if (_assignments[x, y] != newAssignment)
-                        {
-                            _assignments[x, y] = newAssignment;
-                            hasChanged = true;
-                        }
-                    }
-                }
-                RecalculateCentroids();
-            } while (hasChanged);
-
-            stopwatch.Stop();
-
-            UnityEngine.Debug.Log($"Quantization completed in {stopwatch.ElapsedMilliseconds} ms");
 
+            return new QuantizedResult { Indices = indices, Palette = palette };
         }
 
-        private void FillRandomCentroids()
+        private static List<Vector3> KMeans(List<Vector3> colors, int k)
         {
+            List<Vector3> centroids = Enumerable.Range(0, k).Select(i => colors[i * colors.Count / k]).ToList();
 
-            List<Vector3> uniqueColors = new List<Vector3>();
+            List<List<Vector3>> clusters;
-            foreach (Vector3 pixel in _pixels)
+            for (int i = 0; i < 10; i++) // Fixed iterations for performance.... i hate this...
             {
-                if (!uniqueColors.Contains(pixel))
+                clusters = Enumerable.Range(0, k).Select(_ => new List<Vector3>()).ToList();
+                foreach (Vector3 color in colors)
                 {
-                    _uniqueColors.Add(pixel);
+                    int closest = centroids.Select((c, index) => (index, Vector3.SqrMagnitude(c - color)))
+                                           .OrderBy(t => t.Item2).First().index;
+                    clusters[closest].Add(color);
+                }
+
+                for (int j = 0; j < k; j++)
+                {
+                    if (clusters[j].Count > 0)
+                        centroids[j] = clusters[j].Aggregate(Vector3.zero, (acc, c) => acc + c) / clusters[j].Count;
+                }
+            }
+            return centroids;
+        }
+
+        private static void PropagateError(Color[] pixels, int width, int height, int x, int y, Vector3 error)
+        {
+            void AddError(int dx, int dy, float factor)
+            {
+                int nx = x + dx, ny = y + dy;
+                if (nx >= 0 && nx < width && ny >= 0 && ny < height)
+                {
+                    int index = ny * width + nx;
+                    pixels[index].r += error.x * factor;
+                    pixels[index].g += error.y * factor;
+                    pixels[index].b += error.z * factor;
+                }
+            }
+            AddError(1, 0, 7f / 16f);
+            AddError(-1, 1, 3f / 16f);
+            AddError(0, 1, 5f / 16f);
+            AddError(1, 1, 1f / 16f);
+        }
+
+        private static QuantizedResult ConvertToOutput(Color[] pixels, int width, int height)
+        {
+            int[,] indices = new int[width, height];
+            List<Vector3> palette = new List<Vector3>();
+            Dictionary<Vector3, int> colorToIndex = new Dictionary<Vector3, int>();
+
+            for (int y = 0; y < height; y++)
+            {
+                for (int x = 0; x < width; x++)
+                {
+                    Vector3 color = new Vector3(pixels[y * width + x].r, pixels[y * width + x].g, pixels[y * width + x].b);
+                    if (!colorToIndex.ContainsKey(color))
+                    {
+                        colorToIndex[color] = palette.Count;
+                        palette.Add(color);
+                    }
+                    indices[x, y] = colorToIndex[color];
                 }
             }
 
-            for (int i = 0; i < _maxColors; i++)
+            return new QuantizedResult { Indices = indices, Palette = palette };
-            {
-                Vector3 color = _uniqueColors[UnityEngine.Random.Range(0, _uniqueColors.Count - 1)];
-                _centroids[i] = color;
-            }
-
-            kdTree = new KDTree(_centroids);
-
-        }
-
-        private int GetNearestCentroid(Vector3 color)
-        {
-            KDQuery query = new KDQuery();
-            List<int> resultIndices = new List<int>();
-            query.ClosestPoint(kdTree, color, resultIndices);
-            return resultIndices[0];
-        }
-
-        private void RecalculateCentroids()
-        {
-            Vector3[] newCentroids = new Vector3[_maxColors];
-
-
-            for(int i = 0; i < _maxColors; i++) 
-            {
-                List<Vector3> clusterColors = new List<Vector3>();
-                for (int x = 0; x < Width; x++)
-                {
-                    for (int y = 0; y < Height; y++)
-                    {
-                        {
-                            if (_assignments[x, y] == i)
-                            {
-                                clusterColors.Add(_pixels[x, y]);
-                            }
-                        }
         }
     }
 
-                Vector3 newCentroid;
+    public class KDTree
-
-                try
     {
-                    newCentroid = AverageColor(clusterColors);
+        private class Node
-                }
-                catch (InvalidOperationException)
         {
-                    System.Random random = new System.Random();
+            public Vector3 Point;
-                    newCentroid = _uniqueColors[random.Next(0, _uniqueColors.Count - 1)];
+            public Node Left, Right;
         }
 
-                newCentroids[i] = newCentroid;
+        private Node root;
-            }
+        private List<Vector3> points;
 
-            _centroids = newCentroids;
+        public KDTree(List<Vector3> points)
-
-            kdTree = new KDTree(_centroids);
-
-        }
-
-        private Vector3 AverageColor(List<Vector3> colors)
         {
-            float r = colors.Average(c => c.x);
+            this.points = points;
-            float g = colors.Average(c => c.y);
+            root = Build(points, 0);
-            float b = colors.Average(c => c.z);
+        }
-            return new Vector3(r, g, b);
+
+        private Node Build(List<Vector3> points, int depth)
+        {
+            if (points.Count == 0) return null;
+
+            int axis = depth % 3;
+            points.Sort((a, b) => a[axis].CompareTo(b[axis]));
+            int median = points.Count / 2;
+
+            return new Node
+            {
+                Point = points[median],
+                Left = Build(points.Take(median).ToList(), depth + 1),
+                Right = Build(points.Skip(median + 1).ToList(), depth + 1)
+            };
+        }
+
+        public int FindNearestIndex(Vector3 target)
+        {
+            Vector3 nearest = FindNearest(root, target, 0, root.Point);
+            return points.IndexOf(nearest);
+        }
+
+        private Vector3 FindNearest(Node node, Vector3 target, int depth, Vector3 best)
+        {
+            if (node == null) return best;
+
+            if (Vector3.SqrMagnitude(target - node.Point) < Vector3.SqrMagnitude(target - best))
+                best = node.Point;
+
+            int axis = depth % 3;
+            Node first = target[axis] < node.Point[axis] ? node.Left : node.Right;
+            Node second = first == node.Left ? node.Right : node.Left;
+
+            best = FindNearest(first, target, depth + 1, best);
+            if (Mathf.Pow(target[axis] - node.Point[axis], 2) < Vector3.SqrMagnitude(target - best))
+                best = FindNearest(second, target, depth + 1, best);
+
+            return best;
         }
     }
 }

Runtime/PSXTexture2D.cs

@@ -1,5 +1,6 @@
 using System.Collections.Generic;
 using UnityEngine;
+using static PSXSplash.RuntimeCode.TextureQuantizer;
 
 
 namespace PSXSplash.RuntimeCode
@@ -67,7 +68,7 @@ namespace PSXSplash.RuntimeCode
         public ushort[] ImageData { get; set; }
 
 
-        public static PSXTexture2D CreateFromTexture2D(Texture2D inputTexture, PSXBPP bitDepth, bool dither)
+        public static PSXTexture2D CreateFromTexture2D(Texture2D inputTexture, PSXBPP bitDepth)
         {
             PSXTexture2D psxTex = new PSXTexture2D();
 
@@ -90,10 +91,9 @@ namespace PSXSplash.RuntimeCode
 
             psxTex._maxColors = (int)Mathf.Pow((int)bitDepth, 2);
 
-            ImageQuantizer quantizer = new ImageQuantizer();
+            QuantizedResult result = Quantize(inputTexture, psxTex._maxColors);
-            quantizer.Quantize(inputTexture, psxTex._maxColors);
 
-            foreach (Vector3 color in quantizer.Palette)
+            foreach (Vector3 color in result.Palette)
             {
                 Color pixel = new Color(color.x, color.y, color.z);
                 VRAMPixel vramPixel = new VRAMPixel { R = (ushort)(pixel.r * 31), G = (ushort)(pixel.g * 31), B = (ushort)(pixel.b * 31) };
@@ -101,12 +101,12 @@ namespace PSXSplash.RuntimeCode
             }
 
 
-            psxTex.Pixels = new int[quantizer.Width * quantizer.Height];
+            psxTex.Pixels = new int[psxTex.Width * psxTex.Height];
-            for (int x = 0; x < quantizer.Width; x++)
+            for (int x = 0; x < psxTex.Width; x++)
             {
-                for (int y = 0; y < quantizer.Height; y++)
+                for (int y = 0; y < psxTex.Height; y++)
                 {
-                    psxTex.Pixels[x+y*quantizer.Width] = quantizer.Pixels[x,y];
+                    psxTex.Pixels[x + y * psxTex.Width] = result.Indices[x, y];
                 }
             }
 
@@ -217,10 +217,5 @@ namespace PSXSplash.RuntimeCode
             return vramTexture;
 
         }
-
-
-
-
-
     }
 }