Finally doing quantization and dithering in a fast and sensible way

This commit is contained in:
2025-02-11 22:54:17 +01:00
parent ea24f92a2f
commit 0b21a520bc
3 changed files with 138 additions and 148 deletions

View File

@@ -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();

View File

@@ -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 ImageQuantizer
public class TextureQuantizer
{
private int _maxColors;
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
public struct QuantizedResult
{
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();
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 x = 0; x < width; x++)
{
for (int y = 0; y < Height; y++)
{
Color pixel = pixels[x + y * Width];
Vector3 pixelAsVector = new Vector3(pixel.r, pixel.g, pixel.b);
_pixels[x, y] = pixelAsVector;
Vector3 oldColor = new Vector3(pixels[y * width + x].r, pixels[y * width + x].g, pixels[y * width + x].b);
int nearestIndex = kdTree.FindNearestIndex(oldColor);
indices[x, y] = nearestIndex;
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>();
foreach (Vector3 pixel in _pixels)
List<List<Vector3>> clusters;
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++)
{
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]);
}
}
return new QuantizedResult { Indices = indices, Palette = palette };
}
}
Vector3 newCentroid;
try
public class KDTree
{
newCentroid = AverageColor(clusterColors);
}
catch (InvalidOperationException)
private class Node
{
System.Random random = new System.Random();
newCentroid = _uniqueColors[random.Next(0, _uniqueColors.Count - 1)];
public Vector3 Point;
public Node Left, Right;
}
newCentroids[i] = newCentroid;
}
private Node root;
private List<Vector3> points;
_centroids = newCentroids;
kdTree = new KDTree(_centroids);
}
private Vector3 AverageColor(List<Vector3> colors)
public KDTree(List<Vector3> points)
{
float r = colors.Average(c => c.x);
float g = colors.Average(c => c.y);
float b = colors.Average(c => c.z);
return new Vector3(r, g, b);
this.points = points;
root = Build(points, 0);
}
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;
}
}
}

View File

@@ -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();
quantizer.Quantize(inputTexture, psxTex._maxColors);
QuantizedResult result = 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];
for (int x = 0; x < quantizer.Width; x++)
psxTex.Pixels = new int[psxTex.Width * psxTex.Height];
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;
}
}
}