Fix png quantization and reduce Wu memory pressure

src/ImageSharp/Common/Extensions/ComparableExtensions.cs

@@ -137,17 +137,6 @@ namespace SixLabors.ImageSharp
             return value;
         }
 
-        /// <summary>
-        /// Converts an <see cref="int"/> to a <see cref="byte"/> first restricting the value between the
-        /// minimum and maximum allowable ranges.
-        /// </summary>
-        /// <param name="value">The <see cref="int"/> this method extends.</param>
-        /// <returns>The <see cref="byte"/></returns>
-        public static byte ToByte(this int value)
-        {
-            return (byte)value.Clamp(0, 255);
-        }
-
         /// <summary>
         /// Converts an <see cref="float"/> to a <see cref="byte"/> first restricting the value between the
         /// minimum and maximum allowable ranges.

src/ImageSharp/Formats/Png/PngDecoderCore.cs

@@ -217,7 +217,7 @@ namespace SixLabors.ImageSharp.Formats.Png
             where TPixel : struct, IPixel<TPixel>
         {
             var metaData = new ImageMetaData();
-            var pngMetaData = metaData.GetFormatMetaData(PngFormat.Instance);
+            PngMetaData pngMetaData = metaData.GetFormatMetaData(PngFormat.Instance);
             this.currentStream = stream;
             this.currentStream.Skip(8);
             Image<TPixel> image = null;
@@ -307,7 +307,7 @@ namespace SixLabors.ImageSharp.Formats.Png
         public IImageInfo Identify(Stream stream)
         {
             var metaData = new ImageMetaData();
-            var pngMetaData = metaData.GetFormatMetaData(PngFormat.Instance);
+            PngMetaData pngMetaData = metaData.GetFormatMetaData(PngFormat.Instance);
             this.currentStream = stream;
             this.currentStream.Skip(8);
             try

src/ImageSharp/Formats/Png/PngEncoderCore.cs

@@ -243,7 +243,7 @@ namespace SixLabors.ImageSharp.Formats.Png
             // Collect the indexed pixel data
             if (quantized != null)
             {
-                this.WritePaletteChunk(stream, header, quantized);
+                this.WritePaletteChunk(stream, quantized);
             }
 
             this.WritePhysicalChunk(stream, metaData);
@@ -555,30 +555,27 @@ namespace SixLabors.ImageSharp.Formats.Png
         /// </summary>
         /// <typeparam name="TPixel">The pixel format.</typeparam>
         /// <param name="stream">The <see cref="Stream"/> containing image data.</param>
-        /// <param name="header">The <see cref="PngHeader"/>.</param>
         /// <param name="quantized">The quantized frame.</param>
-        private void WritePaletteChunk<TPixel>(Stream stream, in PngHeader header, QuantizedFrame<TPixel> quantized)
+        private void WritePaletteChunk<TPixel>(Stream stream, QuantizedFrame<TPixel> quantized)
             where TPixel : struct, IPixel<TPixel>
         {
             // Grab the palette and write it to the stream.
             TPixel[] palette = quantized.Palette;
-            byte pixelCount = palette.Length.ToByte();
+            int paletteLength = Math.Min(palette.Length, 256);
-
+            int colorTableLength = paletteLength * 3;
-            // Get max colors for bit depth.
-            int colorTableLength = ImageMaths.GetColorCountForBitDepth(header.BitDepth) * 3;
             Rgba32 rgba = default;
             bool anyAlpha = false;
 
             using (IManagedByteBuffer colorTable = this.memoryAllocator.AllocateManagedByteBuffer(colorTableLength))
-            using (IManagedByteBuffer alphaTable = this.memoryAllocator.AllocateManagedByteBuffer(pixelCount))
+            using (IManagedByteBuffer alphaTable = this.memoryAllocator.AllocateManagedByteBuffer(paletteLength))
             {
                 Span<byte> colorTableSpan = colorTable.GetSpan();
                 Span<byte> alphaTableSpan = alphaTable.GetSpan();
                 Span<byte> quantizedSpan = quantized.GetPixelSpan();
 
-                for (byte i = 0; i < pixelCount; i++)
+                for (int i = 0; i < paletteLength; i++)
                 {
-                    if (quantizedSpan.IndexOf(i) > -1)
+                    if (quantizedSpan.IndexOf((byte)i) > -1)
                     {
                         int offset = i * 3;
                         palette[i].ToRgba32(ref rgba);
@@ -604,7 +601,7 @@ namespace SixLabors.ImageSharp.Formats.Png
                 // Write the transparency data
                 if (anyAlpha)
                 {
-                    this.WriteChunk(stream, PngChunkType.PaletteAlpha, alphaTable.Array, 0, pixelCount);
+                    this.WriteChunk(stream, PngChunkType.PaletteAlpha, alphaTable.Array, 0, paletteLength);
                 }
             }
         }

src/ImageSharp/Processing/Processors/Quantization/WuFrameQuantizer{TPixel}.cs

@@ -46,12 +46,12 @@ namespace SixLabors.ImageSharp.Processing.Processors.Quantization
         /// <summary>
         /// The index bits.
         /// </summary>
-        private const int IndexBits = 6;
+        private const int IndexBits = 5;
 
         /// <summary>
         /// The index alpha bits.
         /// </summary>
-        private const int IndexAlphaBits = 3;
+        private const int IndexAlphaBits = 5;
 
         /// <summary>
         /// The index count.
@@ -201,57 +201,10 @@ namespace SixLabors.ImageSharp.Processing.Processors.Quantization
             return this.palette;
         }
 
-        /// <summary>
-        /// Quantizes the pixel
-        /// </summary>
-        /// <param name="rgba">The rgba used to quantize the pixel input</param>
-        private void QuantizePixel(ref Rgba32 rgba)
-        {
-            // Add the color to a 3-D color histogram.
-            int r = rgba.R >> (8 - IndexBits);
-            int g = rgba.G >> (8 - IndexBits);
-            int b = rgba.B >> (8 - IndexBits);
-            int a = rgba.A >> (8 - IndexAlphaBits);
-
-            int index = GetPaletteIndex(r + 1, g + 1, b + 1, a + 1);
-
-            Span<long> vwtSpan = this.vwt.GetSpan();
-            Span<long> vmrSpan = this.vmr.GetSpan();
-            Span<long> vmgSpan = this.vmg.GetSpan();
-            Span<long> vmbSpan = this.vmb.GetSpan();
-            Span<long> vmaSpan = this.vma.GetSpan();
-            Span<float> m2Span = this.m2.GetSpan();
-
-            vwtSpan[index]++;
-            vmrSpan[index] += rgba.R;
-            vmgSpan[index] += rgba.G;
-            vmbSpan[index] += rgba.B;
-            vmaSpan[index] += rgba.A;
-
-            var vector = new Vector4(rgba.R, rgba.G, rgba.B, rgba.A);
-            m2Span[index] += Vector4.Dot(vector, vector);
-        }
-
         /// <inheritdoc/>
         protected override void FirstPass(ImageFrame<TPixel> source, int width, int height)
         {
-            // Build up the 3-D color histogram
+            this.Build3DHistogram(source, width, height);
-            // Loop through each row
-            for (int y = 0; y < height; y++)
-            {
-                Span<TPixel> row = source.GetPixelRowSpan(y);
-                ref TPixel scanBaseRef = ref MemoryMarshal.GetReference(row);
-
-                // And loop through each column
-                Rgba32 rgba = default;
-                for (int x = 0; x < width; x++)
-                {
-                    ref TPixel pixel = ref Unsafe.Add(ref scanBaseRef, x);
-                    pixel.ToRgba32(ref rgba);
-                    this.QuantizePixel(ref rgba);
-                }
-            }
-
             this.Get3DMoments(source.MemoryAllocator);
             this.BuildCube();
         }
@@ -466,6 +419,54 @@ namespace SixLabors.ImageSharp.Processing.Processors.Quantization
             }
         }
 
+        /// <summary>
+        /// Builds a 3-D color histogram of <c>counts, r/g/b, c^2</c>.
+        /// </summary>
+        /// <param name="source">The source data.</param>
+        /// <param name="width">The width in pixels of the image.</param>
+        /// <param name="height">The height in pixels of the image.</param>
+        private void Build3DHistogram(ImageFrame<TPixel> source, int width, int height)
+        {
+            // Build up the 3-D color histogram
+            // Loop through each row
+            Span<long> vwtSpan = this.vwt.GetSpan();
+            Span<long> vmrSpan = this.vmr.GetSpan();
+            Span<long> vmgSpan = this.vmg.GetSpan();
+            Span<long> vmbSpan = this.vmb.GetSpan();
+            Span<long> vmaSpan = this.vma.GetSpan();
+            Span<float> m2Span = this.m2.GetSpan();
+
+            for (int y = 0; y < height; y++)
+            {
+                Span<TPixel> row = source.GetPixelRowSpan(y);
+                ref TPixel scanBaseRef = ref MemoryMarshal.GetReference(row);
+
+                // And loop through each column
+                Rgba32 rgba = default;
+                for (int x = 0; x < width; x++)
+                {
+                    ref TPixel pixel = ref Unsafe.Add(ref scanBaseRef, x);
+                    pixel.ToRgba32(ref rgba);
+
+                    int r = rgba.R >> (8 - IndexBits);
+                    int g = rgba.G >> (8 - IndexBits);
+                    int b = rgba.B >> (8 - IndexBits);
+                    int a = rgba.A >> (8 - IndexAlphaBits);
+
+                    int index = GetPaletteIndex(r + 1, g + 1, b + 1, a + 1);
+
+                    vwtSpan[index]++;
+                    vmrSpan[index] += rgba.R;
+                    vmgSpan[index] += rgba.G;
+                    vmbSpan[index] += rgba.B;
+                    vmaSpan[index] += rgba.A;
+
+                    var vector = new Vector4(rgba.R, rgba.G, rgba.B, rgba.A);
+                    m2Span[index] += Vector4.Dot(vector, vector);
+                }
+            }
+        }
+
         /// <summary>
         /// Converts the histogram into moments so that we can rapidly calculate the sums of the above quantities over any desired box.
         /// </summary>