Add live AV1 luma palette selection

src/ImageSharp/Formats/Heif/Av1/Pipeline/Av1IntraSuperblockEncoder.ModeDecision.cs

@@ -193,6 +193,7 @@ internal static partial class Av1IntraSuperblockEncoder
                 tileIndex,
                 lumaCoefficients[this.codedAreaLuma..],
                 ref lumaState,
+                ref paletteInfo,
                 out int lumaAngleDelta,
                 out Av1FilterIntraMode filterIntraMode);
 
@@ -348,6 +349,7 @@ internal static partial class Av1IntraSuperblockEncoder
             ushort tileIndex,
             Span<int> retainedCoefficients,
             ref Av1EncoderTransformBlockState retainedState,
+            ref Av1EncoderPaletteInfo paletteInfo,
             out int selectedAngleDelta,
             out Av1FilterIntraMode selectedFilterIntraMode)
         {
@@ -463,6 +465,22 @@ internal static partial class Av1IntraSuperblockEncoder
                 BlockSize,
                 TransformSize);
 
+            int paletteDisabledCost = 0;
+            if (this.picture.Parent.FrameHeader.AllowScreenContentTools)
+            {
+                Av1NeighborArrayUnit<Av1EncoderPaletteInfo> paletteContexts = this.picture.PaletteContexts[tileIndex];
+                int blockSizeContext = Av1TileWriter.GetPaletteBlockSizeContext(BlockSize);
+                int neighborContext = Av1TileWriter.GetPaletteYModeContext(
+                    paletteContexts,
+                    macroBlock,
+                    blockOrigin);
+
+                paletteDisabledCost = writer.GetPaletteYModeCost(
+                    false,
+                    blockSizeContext,
+                    neighborContext);
+            }
+
             Span<TSample> candidateReconstruction = stackalloc TSample[SampleCount];
             Span<int> candidateCoefficients = stackalloc int[SampleCount];
             long bestCost = long.MaxValue;
@@ -515,6 +533,7 @@ internal static partial class Av1IntraSuperblockEncoder
                     angleDelta,
                     defaultTransformType,
                     blockContext,
+                    paletteDisabledCost,
                     candidateReconstruction,
                     candidateCoefficients,
                     ref candidateState);
@@ -563,6 +582,7 @@ internal static partial class Av1IntraSuperblockEncoder
                     selectedAngleDelta,
                     transformType,
                     blockContext,
+                    paletteDisabledCost,
                     candidateReconstruction,
                     candidateCoefficients,
                     ref candidateState);
@@ -626,6 +646,7 @@ internal static partial class Av1IntraSuperblockEncoder
                             filterIntraMode,
                             transformType,
                             blockContext,
+                            paletteDisabledCost,
                             candidateReconstruction,
                             candidateCoefficients,
                             ref candidateState);
@@ -651,6 +672,28 @@ internal static partial class Av1IntraSuperblockEncoder
                 }
             }
 
+            if (this.picture.Parent.FrameHeader.AllowScreenContentTools &&
+                this.SelectLumaPalette(
+                    writer,
+                    macroBlock,
+                    sourcePlane,
+                    reconstructionPlane,
+                    blockOrigin,
+                    tileIndex,
+                    transformSetType,
+                    blockContext,
+                    candidateReconstruction,
+                    candidateCoefficients,
+                    retainedCoefficients,
+                    ref retainedState,
+                    ref bestTransformCost,
+                    ref paletteInfo))
+            {
+                bestMode = Av1PredictionMode.DC;
+                selectedAngleDelta = 0;
+                selectedFilterIntraMode = Av1FilterIntraMode.AllFilterIntraModes;
+            }
+
             return bestMode;
         }
 
@@ -667,6 +710,7 @@ internal static partial class Av1IntraSuperblockEncoder
             int angleDelta,
             Av1TransformType transformType,
             Av1TransformBlockContext blockContext,
+            int paletteDisabledCost,
             Span<TSample> candidateReconstruction,
             Span<int> candidateCoefficients,
             ref Av1EncoderTransformBlockState candidateState)
@@ -695,6 +739,11 @@ internal static partial class Av1IntraSuperblockEncoder
                 ref candidateState);
 
             int rate = Av1TileWriter.GetLumaModeCost(writer, macroBlock, BlockSize, mode, angleDelta);
+            if (mode == Av1PredictionMode.DC)
+            {
+                rate += paletteDisabledCost;
+            }
+
             if (mode == Av1PredictionMode.DC && this.picture.Sequence.SequenceHeader.EnableFilterIntra)
             {
                 rate += writer.GetFilterIntraModeCost(Av1FilterIntraMode.AllFilterIntraModes, BlockSize);
@@ -724,6 +773,7 @@ internal static partial class Av1IntraSuperblockEncoder
             Av1FilterIntraMode filterIntraMode,
             Av1TransformType transformType,
             Av1TransformBlockContext blockContext,
+            int paletteDisabledCost,
             Span<TSample> candidateReconstruction,
             Span<int> candidateCoefficients,
             ref Av1EncoderTransformBlockState candidateState)
@@ -754,6 +804,7 @@ internal static partial class Av1IntraSuperblockEncoder
                 Av1PredictionMode.DC,
                 0);
 
+            rate += paletteDisabledCost;
             rate += writer.GetFilterIntraModeCost(filterIntraMode, BlockSize);
             rate += writer.GetCoefficientCost(
                 TransformSize,

src/ImageSharp/Formats/Heif/Av1/Pipeline/Av1IntraSuperblockEncoder.Operator.cs

@@ -2,6 +2,7 @@
 // Licensed under the Six Labors Split License.
 
 using System.Runtime.InteropServices;
+using System.Runtime.Intrinsics;
 using SixLabors.ImageSharp.Formats.Heif.Av1.Prediction;
 using SixLabors.ImageSharp.Formats.Heif.Av1.Prediction.ChromaFromLuma;
 using SixLabors.ImageSharp.Formats.Heif.Av1.Tiling;
@@ -37,6 +38,40 @@ internal static partial class Av1IntraSuperblockEncoder
         /// <returns>The converted sample.</returns>
         public static abstract TSample CreateSample(int value);
 
+        /// <summary>
+        /// Copies active palette-search samples into contiguous signed storage.
+        /// </summary>
+        /// <param name="source">The coded source plane.</param>
+        /// <param name="blockOrigin">The block origin in plane samples.</param>
+        /// <param name="rows">The active row count.</param>
+        /// <param name="columns">The active column count.</param>
+        /// <param name="samples">The contiguous sample destination.</param>
+        public static abstract void CopyPaletteSamples(
+            Buffer2DRegion<TSample> source,
+            Point blockOrigin,
+            int rows,
+            int columns,
+            Span<short> samples);
+
+        /// <summary>
+        /// Builds palette prediction and the matching source residual for transform search.
+        /// </summary>
+        /// <param name="source">The coded source plane.</param>
+        /// <param name="blockOrigin">The block origin in plane samples.</param>
+        /// <param name="paletteColors">The palette colors in index order.</param>
+        /// <param name="colorIndexMap">The complete padded color-index map.</param>
+        /// <param name="prediction">The contiguous prediction destination.</param>
+        /// <param name="residual">The contiguous source-minus-prediction destination.</param>
+        /// <param name="transformSize">The prediction dimensions.</param>
+        public static abstract void PreparePalette(
+            Buffer2DRegion<TSample> source,
+            Point blockOrigin,
+            ReadOnlySpan<ushort> paletteColors,
+            Buffer2DRegion<byte> colorIndexMap,
+            Span<TSample> prediction,
+            Span<short> residual,
+            Av1TransformSize transformSize);
+
         /// <summary>
         /// Builds the zero-mean Q3 luma surface shared by chroma-from-luma candidates.
         /// </summary>
@@ -264,6 +299,71 @@ internal static partial class Av1IntraSuperblockEncoder
         /// <inheritdoc/>
         public static byte CreateSample(int value) => (byte)value;
 
+        /// <inheritdoc/>
+        public static void CopyPaletteSamples(
+            Buffer2DRegion<byte> source,
+            Point blockOrigin,
+            int rows,
+            int columns,
+            Span<short> samples)
+        {
+            int sampleOffset = 0;
+            for (int row = 0; row < rows; row++)
+            {
+                ReadOnlySpan<byte> sourceRow = source
+                    .DangerousGetRowSpan(blockOrigin.Y + row)
+                    .Slice(blockOrigin.X, columns);
+
+                // A complete row widens in one vector; clipped edge rows retain scalar bounds.
+                if (columns == 8 && Vector128.IsHardwareAccelerated)
+                {
+                    ulong packed = MemoryMarshal.Read<ulong>(sourceRow);
+                    Vector128.WidenLower(Vector128.CreateScalarUnsafe(packed).AsByte())
+                        .AsInt16()
+                        .CopyTo(samples[sampleOffset..]);
+
+                    sampleOffset += columns;
+                    continue;
+                }
+
+                for (int column = 0; column < sourceRow.Length; column++)
+                {
+                    samples[sampleOffset++] = sourceRow[column];
+                }
+            }
+        }
+
+        /// <inheritdoc/>
+        public static void PreparePalette(
+            Buffer2DRegion<byte> source,
+            Point blockOrigin,
+            ReadOnlySpan<ushort> paletteColors,
+            Buffer2DRegion<byte> colorIndexMap,
+            Span<byte> prediction,
+            Span<short> residual,
+            Av1TransformSize transformSize)
+        {
+            int width = transformSize.GetWidth();
+            int height = transformSize.GetHeight();
+            Av1PalettePredictor.Predict(
+                paletteColors,
+                colorIndexMap,
+                prediction,
+                width,
+                width,
+                height);
+
+            Av1ResidualBuilder.Subtract(
+                Av1TransformBlockEncoder.GetPlaneSpan(source, blockOrigin),
+                source.Stride,
+                prediction,
+                width,
+                residual,
+                width,
+                width,
+                height);
+        }
+
         /// <inheritdoc/>
         public static void PrepareChromaFromLuma(
             Buffer2DRegion<byte> reconstruction,
@@ -493,6 +593,57 @@ internal static partial class Av1IntraSuperblockEncoder
         /// <inheritdoc/>
         public static ushort CreateSample(int value) => (ushort)value;
 
+        /// <inheritdoc/>
+        public static void CopyPaletteSamples(
+            Buffer2DRegion<ushort> source,
+            Point blockOrigin,
+            int rows,
+            int columns,
+            Span<short> samples)
+        {
+            int sampleOffset = 0;
+            for (int row = 0; row < rows; row++)
+            {
+                ReadOnlySpan<ushort> sourceRow = source
+                    .DangerousGetRowSpan(blockOrigin.Y + row)
+                    .Slice(blockOrigin.X, columns);
+
+                MemoryMarshal.Cast<ushort, short>(sourceRow).CopyTo(samples[sampleOffset..]);
+                sampleOffset += sourceRow.Length;
+            }
+        }
+
+        /// <inheritdoc/>
+        public static void PreparePalette(
+            Buffer2DRegion<ushort> source,
+            Point blockOrigin,
+            ReadOnlySpan<ushort> paletteColors,
+            Buffer2DRegion<byte> colorIndexMap,
+            Span<ushort> prediction,
+            Span<short> residual,
+            Av1TransformSize transformSize)
+        {
+            int width = transformSize.GetWidth();
+            int height = transformSize.GetHeight();
+            Av1PalettePredictor.Predict(
+                paletteColors,
+                colorIndexMap,
+                MemoryMarshal.Cast<ushort, short>(prediction),
+                width,
+                width,
+                height);
+
+            Av1ResidualBuilder.Subtract(
+                Av1TransformBlockEncoder.GetPlaneSpan(source, blockOrigin),
+                source.Stride,
+                prediction,
+                width,
+                residual,
+                width,
+                width,
+                height);
+        }
+
         /// <inheritdoc/>
         public static void PrepareChromaFromLuma(
             Buffer2DRegion<ushort> reconstruction,

src/ImageSharp/Formats/Heif/Av1/Pipeline/Av1IntraSuperblockEncoder.PaletteModeDecision.cs

@@ -0,0 +1,416 @@
+// Copyright (c) Six Labors.
+// Licensed under the Six Labors Split License.
+
+using SixLabors.ImageSharp.Formats.Heif.Av1.Entropy;
+using SixLabors.ImageSharp.Formats.Heif.Av1.OpenBitstreamUnit;
+using SixLabors.ImageSharp.Formats.Heif.Av1.Prediction;
+using SixLabors.ImageSharp.Formats.Heif.Av1.Tiling;
+using SixLabors.ImageSharp.Formats.Heif.Av1.Transform;
+using SixLabors.ImageSharp.Memory;
+
+namespace SixLabors.ImageSharp.Formats.Heif.Av1.Pipeline;
+
+/// <content>
+/// Provides luma palette mode decisions for intra encoding.
+/// </content>
+internal static partial class Av1IntraSuperblockEncoder
+{
+    internal partial struct ModeDecision<TSample, TOperator>
+        where TSample : unmanaged
+        where TOperator : struct, IBlockEncodingOperator<TSample>
+    {
+        private bool SelectLumaPalette(
+            Av1SymbolEncoder writer,
+            Av1MacroBlockD macroBlock,
+            Buffer2DRegion<TSample> sourcePlane,
+            Buffer2DRegion<TSample> reconstructionPlane,
+            Point blockOrigin,
+            ushort tileIndex,
+            Av1TransformSetType transformSetType,
+            Av1TransformBlockContext blockContext,
+            Span<TSample> candidateReconstruction,
+            Span<int> candidateCoefficients,
+            Span<int> retainedCoefficients,
+            ref Av1EncoderTransformBlockState retainedState,
+            ref long bestCost,
+            ref Av1EncoderPaletteInfo paletteInfo)
+        {
+            const Av1BlockSize BlockSize = Av1BlockSize.Block8x8;
+            const int BlockLength = 8;
+            const int SampleCapacity = BlockLength * BlockLength;
+            ObuFrameSize frameSize = this.picture.Parent.FrameHeader.FrameSize;
+            int rows = Math.Min(BlockLength, frameSize.FrameHeight - blockOrigin.Y);
+            int columns = Math.Min(BlockLength, frameSize.FrameWidth - blockOrigin.X);
+            int sampleCount = rows * columns;
+            Span<short> samples = stackalloc short[SampleCapacity];
+            samples = samples[..sampleCount];
+            TOperator.CopyPaletteSamples(sourcePlane, blockOrigin, rows, columns, samples);
+
+            Span<short> uniqueColors = stackalloc short[SampleCapacity];
+            Span<int> colorCounts = stackalloc int[SampleCapacity];
+            int uniqueColorCount = 0;
+            short minimum = samples[0];
+            short maximum = samples[0];
+            foreach (short sample in samples)
+            {
+                int colorIndex = uniqueColors[..uniqueColorCount].IndexOf(sample);
+                if (colorIndex >= 0)
+                {
+                    colorCounts[colorIndex]++;
+                }
+                else
+                {
+                    uniqueColors[uniqueColorCount] = sample;
+                    colorCounts[uniqueColorCount] = 1;
+                    uniqueColorCount++;
+                }
+
+                minimum = Math.Min(minimum, sample);
+                maximum = Math.Max(maximum, sample);
+            }
+
+            if (uniqueColorCount < 2)
+            {
+                return false;
+            }
+
+            int maximumPaletteSize = Math.Min(uniqueColorCount, Av1Constants.PaletteMaxSize);
+            Span<byte> dominantOrder = stackalloc byte[SampleCapacity];
+            for (int index = 0; index < uniqueColorCount; index++)
+            {
+                dominantOrder[index] = (byte)index;
+            }
+
+            // Count order chooses the colors that explain most samples first; sample value resolves equal counts.
+            for (int index = 1; index < uniqueColorCount; index++)
+            {
+                byte current = dominantOrder[index];
+                int destination = index;
+                while (destination > 0)
+                {
+                    byte preceding = dominantOrder[destination - 1];
+                    bool precedes = colorCounts[current] > colorCounts[preceding] ||
+                        (colorCounts[current] == colorCounts[preceding] && uniqueColors[current] < uniqueColors[preceding]);
+
+                    if (!precedes)
+                    {
+                        break;
+                    }
+
+                    dominantOrder[destination] = preceding;
+                    destination--;
+                }
+
+                dominantOrder[destination] = current;
+            }
+
+            Av1NeighborArrayUnit<Av1EncoderPaletteInfo> paletteContexts = this.picture.PaletteContexts[tileIndex];
+            int blockSizeContext = Av1TileWriter.GetPaletteBlockSizeContext(BlockSize);
+            int neighborContext = Av1TileWriter.GetPaletteYModeContext(paletteContexts, macroBlock, blockOrigin);
+            Span<ushort> colorCache = stackalloc ushort[2 * Av1Constants.PaletteMaxSize];
+            int colorCacheSize = Av1TileWriter.GetPaletteCache(
+                paletteContexts,
+                macroBlock,
+                blockOrigin,
+                Av1Plane.Y,
+                colorCache);
+
+            Buffer2DRegion<byte> colorIndexMap = this.superblock.Workspace
+                .GetPaletteMaps()
+                .GetMap(Av1PlaneType.Y, BlockLength, BlockLength);
+
+            Span<byte> retainedColorIndexMap = stackalloc byte[SampleCapacity];
+            Span<short> centroids = stackalloc short[Av1Constants.PaletteMaxSize];
+            bool paletteSelected = false;
+
+            // Exhaustive ascending size search avoids the reference encoder's speed-dependent pruning.
+            for (int paletteSize = 2; paletteSize <= maximumPaletteSize; paletteSize++)
+            {
+                for (int index = 0; index < paletteSize; index++)
+                {
+                    centroids[index] = uniqueColors[dominantOrder[index]];
+                }
+
+                this.EvaluateLumaPaletteCandidate(
+                    writer,
+                    macroBlock,
+                    blockOrigin,
+                    transformSetType,
+                    blockContext,
+                    samples,
+                    rows,
+                    columns,
+                    colorCache[..colorCacheSize],
+                    blockSizeContext,
+                    neighborContext,
+                    centroids[..paletteSize],
+                    colorIndexMap,
+                    candidateReconstruction,
+                    candidateCoefficients,
+                    retainedCoefficients,
+                    retainedColorIndexMap,
+                    reconstructionPlane,
+                    ref retainedState,
+                    ref bestCost,
+                    ref paletteInfo,
+                    ref paletteSelected);
+            }
+
+            if (uniqueColorCount == 2)
+            {
+                centroids[0] = minimum;
+                centroids[1] = maximum;
+                this.EvaluateLumaPaletteCandidate(
+                    writer,
+                    macroBlock,
+                    blockOrigin,
+                    transformSetType,
+                    blockContext,
+                    samples,
+                    rows,
+                    columns,
+                    colorCache[..colorCacheSize],
+                    blockSizeContext,
+                    neighborContext,
+                    centroids[..2],
+                    colorIndexMap,
+                    candidateReconstruction,
+                    candidateCoefficients,
+                    retainedCoefficients,
+                    retainedColorIndexMap,
+                    reconstructionPlane,
+                    ref retainedState,
+                    ref bestCost,
+                    ref paletteInfo,
+                    ref paletteSelected);
+            }
+            else
+            {
+                Span<byte> clusterIndices = stackalloc byte[SampleCapacity];
+                clusterIndices = clusterIndices[..sampleCount];
+                for (int paletteSize = 2; paletteSize <= maximumPaletteSize; paletteSize++)
+                {
+                    Span<short> candidateCentroids = centroids[..paletteSize];
+                    Av1PaletteKMeans.InitializeCentroids(minimum, maximum, candidateCentroids);
+                    Av1PaletteKMeans.Cluster(samples, candidateCentroids, clusterIndices);
+                    this.EvaluateLumaPaletteCandidate(
+                        writer,
+                        macroBlock,
+                        blockOrigin,
+                        transformSetType,
+                        blockContext,
+                        samples,
+                        rows,
+                        columns,
+                        colorCache[..colorCacheSize],
+                        blockSizeContext,
+                        neighborContext,
+                        candidateCentroids,
+                        colorIndexMap,
+                        candidateReconstruction,
+                        candidateCoefficients,
+                        retainedCoefficients,
+                        retainedColorIndexMap,
+                        reconstructionPlane,
+                        ref retainedState,
+                        ref bestCost,
+                        ref paletteInfo,
+                        ref paletteSelected);
+                }
+            }
+
+            if (paletteSelected)
+            {
+                for (int row = 0; row < BlockLength; row++)
+                {
+                    retainedColorIndexMap.Slice(row * BlockLength, BlockLength)
+                        .CopyTo(colorIndexMap.DangerousGetRowSpan(row));
+                }
+            }
+
+            return paletteSelected;
+        }
+
+        private void EvaluateLumaPaletteCandidate(
+            Av1SymbolEncoder writer,
+            Av1MacroBlockD macroBlock,
+            Point blockOrigin,
+            Av1TransformSetType transformSetType,
+            Av1TransformBlockContext blockContext,
+            ReadOnlySpan<short> samples,
+            int rows,
+            int columns,
+            ReadOnlySpan<ushort> colorCache,
+            int blockSizeContext,
+            int neighborContext,
+            Span<short> centroids,
+            Buffer2DRegion<byte> colorIndexMap,
+            Span<TSample> candidateReconstruction,
+            Span<int> candidateCoefficients,
+            Span<int> retainedCoefficients,
+            Span<byte> retainedColorIndexMap,
+            Buffer2DRegion<TSample> reconstructionPlane,
+            ref Av1EncoderTransformBlockState retainedState,
+            ref long bestCost,
+            ref Av1EncoderPaletteInfo paletteInfo,
+            ref bool paletteSelected)
+        {
+            const Av1BlockSize BlockSize = Av1BlockSize.Block8x8;
+            const Av1TransformSize TransformSize = Av1TransformSize.Size8x8;
+            const int BlockLength = 8;
+            const int SampleCount = BlockLength * BlockLength;
+            int bitDepth = this.bitDepth.GetBitCount();
+            int cacheThreshold = 4 << (bitDepth - 8);
+            for (int colorIndex = 0; colorIndex < centroids.Length && !colorCache.IsEmpty; colorIndex++)
+            {
+                int minimumDifference = Math.Abs(centroids[colorIndex] - colorCache[0]);
+                int nearestCacheIndex = 0;
+                for (int cacheIndex = 1; cacheIndex < colorCache.Length; cacheIndex++)
+                {
+                    int difference = Math.Abs(centroids[colorIndex] - colorCache[cacheIndex]);
+                    if (difference < minimumDifference)
+                    {
+                        minimumDifference = difference;
+                        nearestCacheIndex = cacheIndex;
+                    }
+                }
+
+                if (minimumDifference <= cacheThreshold)
+                {
+                    centroids[colorIndex] = (short)colorCache[nearestCacheIndex];
+                }
+            }
+
+            centroids.Sort();
+            int paletteSize = 1;
+            for (int colorIndex = 1; colorIndex < centroids.Length; colorIndex++)
+            {
+                if (centroids[colorIndex] != centroids[colorIndex - 1])
+                {
+                    centroids[paletteSize++] = centroids[colorIndex];
+                }
+            }
+
+            if (paletteSize < 2)
+            {
+                return;
+            }
+
+            ReadOnlySpan<short> paletteCentroids = centroids[..paletteSize];
+            Span<ushort> paletteColors = stackalloc ushort[Av1Constants.PaletteMaxSize];
+            paletteColors = paletteColors[..paletteSize];
+            for (int colorIndex = 0; colorIndex < paletteSize; colorIndex++)
+            {
+                paletteColors[colorIndex] = (ushort)paletteCentroids[colorIndex];
+            }
+
+            Span<byte> colorIndices = stackalloc byte[SampleCount];
+            Av1PaletteKMeans.AssignIndices(samples, paletteCentroids, colorIndices);
+            for (int row = 0; row < rows; row++)
+            {
+                Span<byte> mapRow = colorIndexMap.DangerousGetRowSpan(row)[..BlockLength];
+                colorIndices.Slice(row * columns, columns).CopyTo(mapRow);
+                mapRow[columns..].Fill(mapRow[columns - 1]);
+            }
+
+            // Padding repeats the last active edge so transform prediction matches coded-frame edge extension.
+            for (int row = rows; row < BlockLength; row++)
+            {
+                colorIndexMap.DangerousGetRowSpan(rows - 1)[..BlockLength]
+                    .CopyTo(colorIndexMap.DangerousGetRowSpan(row));
+            }
+
+            Span<TSample> prediction = stackalloc TSample[SampleCount];
+            Span<short> residual = stackalloc short[SampleCount];
+            TOperator.PreparePalette(
+                this.source.GetPlane(Av1Plane.Y),
+                blockOrigin,
+                paletteColors,
+                colorIndexMap,
+                prediction,
+                residual,
+                TransformSize);
+
+            int rate = Av1TileWriter.GetLumaModeCost(
+                writer,
+                macroBlock,
+                BlockSize,
+                Av1PredictionMode.DC,
+                0);
+
+            rate += writer.GetPaletteYModeCost(true, blockSizeContext, neighborContext);
+            rate += writer.GetPaletteSizeCost(paletteSize, blockSizeContext, Av1PlaneType.Y);
+            rate += Av1SymbolEncoder.GetPaletteYColorCost(colorCache, paletteColors, bitDepth);
+            rate += writer.GetPaletteColorMapCost(
+                paletteSize,
+                Av1PlaneType.Y,
+                rows,
+                columns,
+                colorIndexMap);
+
+            for (Av1TransformType transformType = Av1TransformType.DctDct;
+                transformType < Av1TransformType.AllTransformTypes;
+                transformType++)
+            {
+                if (!transformType.IsExtendedSetUsed(transformSetType))
+                {
+                    continue;
+                }
+
+                Av1EncoderTransformBlockState candidateState = default;
+                long distortion = TOperator.EncodePredictionCandidate(
+                    this.blockWorkspace,
+                    this.source.GetPlane(Av1Plane.Y),
+                    blockOrigin,
+                    prediction,
+                    residual,
+                    candidateReconstruction,
+                    candidateCoefficients,
+                    TransformSize,
+                    transformType,
+                    Av1Plane.Y,
+                    this.quantization.QIndex[0],
+                    this.quantization.DeltaQDc[(int)Av1Plane.Y],
+                    this.quantization.DeltaQAc[(int)Av1Plane.Y],
+                    this.bitDepth,
+                    ref candidateState);
+
+                int candidateRate = rate + writer.GetCoefficientCost(
+                    TransformSize,
+                    transformType,
+                    Av1PredictionMode.DC,
+                    candidateCoefficients,
+                    Av1ComponentType.Luminance,
+                    blockContext,
+                    candidateState.EndOfBlock,
+                    this.picture.Parent.FrameHeader.UseReducedTransformSet,
+                    Av1FilterIntraMode.AllFilterIntraModes);
+
+                long candidateCost = Av1RateDistortion.GetCost(this.rateMultiplier, candidateRate, distortion);
+                if (candidateCost < bestCost)
+                {
+                    CopyCandidate(
+                        candidateReconstruction,
+                        candidateCoefficients,
+                        reconstructionPlane,
+                        blockOrigin,
+                        retainedCoefficients,
+                        TransformSize,
+                        candidateState,
+                        ref retainedState);
+
+                    for (int row = 0; row < BlockLength; row++)
+                    {
+                        colorIndexMap.DangerousGetRowSpan(row)[..BlockLength]
+                            .CopyTo(retainedColorIndexMap[(row * BlockLength)..]);
+                    }
+
+                    paletteInfo.PaletteSizes[0] = (byte)paletteSize;
+                    paletteInfo.SetColors(Av1Plane.Y, paletteColors);
+                    bestCost = candidateCost;
+                    paletteSelected = true;
+                }
+            }
+        }
+    }
+}

src/ImageSharp/Formats/Heif/Av1/Tiling/Av1EncoderSuperblockWorkspace.cs

@@ -59,7 +59,7 @@ internal sealed class Av1EncoderSuperblockWorkspace : IDisposable
     public ref Av1EncoderPaletteInfo PaletteInfo => ref this.paletteInfo;
 
     /// <summary>
-    /// Gets the reusable palette maps, allocating their shared owner only after a block selects palette mode.
+    /// Gets the reusable palette maps, allocating their shared owner only after a block enters palette search.
     /// </summary>
     /// <returns>The reusable luma and chroma palette maps.</returns>
     public Av1EncoderPaletteMapBuffer GetPaletteMaps()

src/ImageSharp/Formats/Heif/Av1/Tiling/Av1TileWriter.cs

@@ -1157,24 +1157,12 @@ internal partial class Av1TileWriter
         int tileIndex,
         bool hasChroma)
     {
-        int blockSizeContext = Av1Math.Log2(blockSize.GetWidth() * blockSize.GetHeight()) - 6;
+        int blockSizeContext = GetPaletteBlockSizeContext(blockSize);
         Av1NeighborArrayUnit<Av1EncoderPaletteInfo> paletteContexts = pcs.PaletteContexts[tileIndex];
         int yPaletteSize = paletteInfo.PaletteSizes[0];
         if (macroBlockModeInfo.Block.Mode == Av1PredictionMode.DC)
         {
-            int neighborContext = 0;
+            int neighborContext = GetPaletteYModeContext(paletteContexts, macroBlock, blockOrigin);
-            if (macroBlock.IsUpAvailable &&
-                paletteContexts.Top[paletteContexts.GetTopIndex(blockOrigin)].PaletteSizes[0] != 0)
-            {
-                neighborContext++;
-            }
-
-            if (macroBlock.IsLeftAvailable &&
-                paletteContexts.Left[paletteContexts.GetLeftIndex(blockOrigin)].PaletteSizes[0] != 0)
-            {
-                neighborContext++;
-            }
-
             writer.WritePaletteYMode(yPaletteSize != 0, blockSizeContext, neighborContext);
             if (yPaletteSize != 0)
             {
@@ -1223,7 +1211,7 @@ internal partial class Av1TileWriter
     /// <summary>
     /// Builds the sorted palette-color cache from the available above and left encoder edges.
     /// </summary>
-    private static int GetPaletteCache(
+    internal static int GetPaletteCache(
         Av1NeighborArrayUnit<Av1EncoderPaletteInfo> paletteContexts,
         Av1MacroBlockD macroBlock,
         Point blockOrigin,
@@ -1247,6 +1235,36 @@ internal partial class Av1TileWriter
         return Av1PaletteCache.Merge(aboveColors, leftColors, cache);
     }
 
+    /// <summary>
+    /// Gets the palette probability context derived from the logarithmic block area.
+    /// </summary>
+    internal static int GetPaletteBlockSizeContext(Av1BlockSize blockSize)
+        => Av1Math.Log2(blockSize.GetWidth() * blockSize.GetHeight()) - 6;
+
+    /// <summary>
+    /// Counts the available above and left luma neighbors that selected palette mode.
+    /// </summary>
+    internal static int GetPaletteYModeContext(
+        Av1NeighborArrayUnit<Av1EncoderPaletteInfo> paletteContexts,
+        Av1MacroBlockD macroBlock,
+        Point blockOrigin)
+    {
+        int neighborContext = 0;
+        if (macroBlock.IsUpAvailable &&
+            paletteContexts.Top[paletteContexts.GetTopIndex(blockOrigin)].PaletteSizes[0] != 0)
+        {
+            neighborContext++;
+        }
+
+        if (macroBlock.IsLeftAvailable &&
+            paletteContexts.Left[paletteContexts.GetLeftIndex(blockOrigin)].PaletteSizes[0] != 0)
+        {
+            neighborContext++;
+        }
+
+        return neighborContext;
+    }
+
     /// <summary>
     /// Determines whether filter-intra syntax is available for a block mode.
     /// </summary>

tests/ImageSharp.Tests/Formats/Heif/Av1/Av1IntraSuperblockEncoderTests.cs

@@ -716,6 +716,70 @@ public class Av1IntraSuperblockEncoderTests
         Assert.False(payloads[0].SequenceEqual(payloads[1]));
     }
 
+    [Fact]
+    public void ProductionTileSelectsExactLumaPaletteAtFullAndClippedSizes()
+    {
+        AssertProductionTileSelectsExactLumaPalette(
+            Av1BitDepth.EightBit,
+            8,
+            8,
+            8,
+            (byte)32,
+            (byte)224,
+            32,
+            224,
+            static (source, reconstruction, picture, coefficients, superblockWorkspace, blockWorkspace) =>
+                new Av1IntraTileWriter(
+                    Configuration.Default,
+                    source,
+                    reconstruction,
+                    picture,
+                    coefficients,
+                    superblockWorkspace,
+                    blockWorkspace,
+                    initialSize: 256));
+
+        AssertProductionTileSelectsExactLumaPalette(
+            Av1BitDepth.TwelveBit,
+            12,
+            8,
+            8,
+            (ushort)512,
+            (ushort)3584,
+            512,
+            3584,
+            static (source, reconstruction, picture, coefficients, superblockWorkspace, blockWorkspace) =>
+                new Av1IntraTileWriter(
+                    Configuration.Default,
+                    source,
+                    reconstruction,
+                    picture,
+                    coefficients,
+                    superblockWorkspace,
+                    blockWorkspace,
+                    initialSize: 256));
+
+        AssertProductionTileSelectsExactLumaPalette(
+            Av1BitDepth.EightBit,
+            8,
+            5,
+            3,
+            (byte)48,
+            (byte)208,
+            48,
+            208,
+            static (source, reconstruction, picture, coefficients, superblockWorkspace, blockWorkspace) =>
+                new Av1IntraTileWriter(
+                    Configuration.Default,
+                    source,
+                    reconstruction,
+                    picture,
+                    coefficients,
+                    superblockWorkspace,
+                    blockWorkspace,
+                    initialSize: 256));
+    }
+
     [Theory]
     [InlineData((int)Av1PredictionMode.Vertical, 0)]
     [InlineData((int)Av1PredictionMode.Horizontal, 0)]
@@ -1835,6 +1899,120 @@ public class Av1IntraSuperblockEncoderTests
         };
     }
 
+    private static void AssertProductionTileSelectsExactLumaPalette<TSample>(
+        Av1BitDepth bitDepth,
+        int bitDepthValue,
+        int width,
+        int height,
+        TSample lowerColor,
+        TSample upperColor,
+        ushort expectedLowerColor,
+        ushort expectedUpperColor,
+        TileWriterFactory<TSample> createTileWriter)
+        where TSample : unmanaged
+    {
+        const int QIndex = 37;
+        ObuColorConfig colorConfig = new()
+        {
+            IsMonochrome = true,
+            SubSamplingX = true,
+            SubSamplingY = true,
+            BitDepth = bitDepth
+        };
+
+        using Av1EncoderFrameBuffer<TSample> source = new(
+            Configuration.Default,
+            width,
+            height,
+            bitDepthValue,
+            Av1ColorFormat.Yuv400,
+            0,
+            0);
+
+        using Av1EncoderFrameBuffer<TSample> reconstruction = new(
+            Configuration.Default,
+            width,
+            height,
+            bitDepthValue,
+            Av1ColorFormat.Yuv400,
+            0,
+            0);
+
+        Buffer2DRegion<TSample> sourcePlane = source.Frame.CodedView.GetPlane(Av1Plane.Y);
+        for (int row = 0; row < sourcePlane.Height; row++)
+        {
+            int visibleRow = Math.Min(row, height - 1);
+            sourcePlane.DangerousGetRowSpan(row).Fill(visibleRow < height / 2 ? lowerColor : upperColor);
+        }
+
+        ClearPlane(reconstruction.Luma);
+        using Av1EncoderModeInfoBuffer modeInfo = new(
+            Configuration.Default,
+            width,
+            height,
+            disallow4x4AllFrames: true);
+
+        Av1PictureControlSet pictureTemplate = CreatePicture(
+            modeInfo,
+            colorConfig,
+            use128x128Superblock: false,
+            QIndex);
+
+        pictureTemplate.Parent.FrameHeader.AllowScreenContentTools = true;
+        pictureTemplate.Parent.FrameHeader.FrameSize.FrameWidth = width;
+        pictureTemplate.Parent.FrameHeader.FrameSize.FrameHeight = height;
+        using Av1EncoderPictureBuffer picture = new(
+            Configuration.Default,
+            pictureTemplate.Sequence.SequenceHeader,
+            pictureTemplate.Parent.FrameHeader,
+            width,
+            height);
+
+        using Av1EncoderCoefficientBuffer coefficients = new(
+            Configuration.Default,
+            pictureTemplate.Sequence.SequenceHeader,
+            width,
+            height);
+
+        using Av1EncoderSuperblockWorkspace superblockWorkspace = new(Configuration.Default);
+        using Av1EncoderBlockWorkspace blockWorkspace = new(Configuration.Default);
+        using Av1IntraTileWriter tileWriter = createTileWriter(
+            source.Frame,
+            reconstruction.Frame,
+            picture.Picture,
+            coefficients,
+            superblockWorkspace,
+            blockWorkspace);
+
+        ref Av1MacroBlockModeInfo mode = ref picture.Picture.GetMacroBlockModeInfo(default);
+        Assert.Equal(Av1PredictionMode.DC, mode.Block.Mode);
+        Assert.Equal(Av1FilterIntraMode.AllFilterIntraModes, superblockWorkspace.FinalBlocks[0].FilterIntraMode);
+        Assert.Equal((ushort)0, coefficients.GetTransformBlockSpan(0, Av1Plane.Y)[0].EndOfBlock);
+        Assert.Equal(2, superblockWorkspace.PaletteInfo.PaletteSizes[0]);
+        Assert.Equal(
+            [expectedLowerColor, expectedUpperColor],
+            superblockWorkspace.PaletteInfo.GetColors(Av1Plane.Y).ToArray());
+
+        Buffer2DRegion<byte> colorIndexMap = superblockWorkspace
+            .GetPaletteMaps()
+            .GetMap(Av1PlaneType.Y, 8, 8);
+
+        Buffer2DRegion<TSample> reconstructionPlane = reconstruction.Frame.CodedView.GetPlane(Av1Plane.Y);
+        for (int row = 0; row < reconstructionPlane.Height; row++)
+        {
+            int visibleRow = Math.Min(row, height - 1);
+            byte expectedIndex = (byte)(visibleRow < height / 2 ? 0 : 1);
+            foreach (byte index in colorIndexMap.DangerousGetRowSpan(row))
+            {
+                Assert.Equal(expectedIndex, index);
+            }
+
+            Assert.True(sourcePlane.DangerousGetRowSpan(row).SequenceEqual(reconstructionPlane.DangerousGetRowSpan(row)));
+        }
+
+        Assert.NotEqual(0, tileWriter.GetTileData(0).Length);
+    }
+
     private static void FillChromaModeSelectionPlane(
         Buffer2DRegion<byte> plane,
         Av1TransformSize transformSize,