Browse Source

Add live AV1 chroma palette selection

pull/2633/head
James Jackson-South 1 month ago
parent
commit
d17b8d08d9
  1. 5
      HEIF_IMPLEMENTATION_PLAN.md
  2. 41
      src/ImageSharp/Formats/Heif/Av1/Pipeline/Av1IntraSuperblockEncoder.ChromaModeDecision.cs
  3. 373
      src/ImageSharp/Formats/Heif/Av1/Pipeline/Av1IntraSuperblockEncoder.ChromaPaletteModeDecision.cs
  4. 1
      src/ImageSharp/Formats/Heif/Av1/Pipeline/Av1IntraSuperblockEncoder.ModeDecision.cs
  5. 133
      tests/ImageSharp.Tests/Formats/Heif/Av1/Av1IntraSuperblockEncoderTests.cs

5
HEIF_IMPLEMENTATION_PLAN.md

File diff suppressed because one or more lines are too long

41
src/ImageSharp/Formats/Heif/Av1/Pipeline/Av1IntraSuperblockEncoder.ChromaModeDecision.cs

@ -53,6 +53,7 @@ internal static partial class Av1IntraSuperblockEncoder
Span<int> retainedRedCoefficients,
ref Av1EncoderTransformBlockState retainedBlueState,
ref Av1EncoderTransformBlockState retainedRedState,
ref Av1EncoderPaletteInfo paletteInfo,
out int selectedAngleDelta,
out byte selectedChromaFromLumaIndex,
out sbyte selectedChromaFromLumaSigns)
@ -163,6 +164,10 @@ internal static partial class Av1IntraSuperblockEncoder
int deltaCount = AngleDeltaSearchOrder.Length;
int directionalModeCount = (int)Av1ChromaPredictionMode.Directional67Degrees - (int)Av1ChromaPredictionMode.Vertical + 1;
int candidateCount = baseModeCount + (directionalModeCount * deltaCount);
bool hasLumaPalette = paletteInfo.PaletteSizes[0] != 0;
int paletteDisabledCost = this.picture.Parent.FrameHeader.AllowScreenContentTools
? writer.GetPaletteUvModeCost(false, hasLumaPalette)
: 0;
// Spatial base modes precede the six nonzero adjustments for each directional mode.
// Chroma-from-luma remains a separate search because it consumes reconstructed luma AC state.
@ -202,6 +207,7 @@ internal static partial class Av1IntraSuperblockEncoder
hasAbove,
blueContext,
redContext,
paletteDisabledCost,
candidateBlueReconstruction[..sampleCount],
candidateRedReconstruction[..sampleCount],
candidateBlueCoefficients[..sampleCount],
@ -434,6 +440,35 @@ internal static partial class Av1IntraSuperblockEncoder
}
}
if (this.picture.Parent.FrameHeader.AllowScreenContentTools &&
this.SelectChromaPalette(
writer,
macroBlock,
modeInfo,
lumaOrigin,
chromaOrigin,
tileIndex,
lumaMode,
transformSize,
blueContext,
redContext,
candidateBlueReconstruction[..sampleCount],
candidateRedReconstruction[..sampleCount],
candidateBlueCoefficients[..sampleCount],
candidateRedCoefficients[..sampleCount],
retainedBlueCoefficients,
retainedRedCoefficients,
ref retainedBlueState,
ref retainedRedState,
ref bestCost,
ref paletteInfo))
{
bestMode = Av1ChromaPredictionMode.DC;
selectedAngleDelta = 0;
selectedChromaFromLumaIndex = 0;
selectedChromaFromLumaSigns = 0;
}
return bestMode;
}
@ -502,6 +537,7 @@ internal static partial class Av1IntraSuperblockEncoder
bool hasAbove,
Av1TransformBlockContext blueContext,
Av1TransformBlockContext redContext,
int paletteDisabledCost,
Span<TSample> candidateBlueReconstruction,
Span<TSample> candidateRedReconstruction,
Span<int> candidateBlueCoefficients,
@ -568,6 +604,11 @@ internal static partial class Av1IntraSuperblockEncoder
chromaMode,
angleDelta);
if (chromaMode == Av1ChromaPredictionMode.DC)
{
rate += paletteDisabledCost;
}
rate += writer.GetCoefficientCost(
transformSize,
transformType,

373
src/ImageSharp/Formats/Heif/Av1/Pipeline/Av1IntraSuperblockEncoder.ChromaPaletteModeDecision.cs

@ -0,0 +1,373 @@
// 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 paired chroma 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 SelectChromaPalette(
Av1SymbolEncoder writer,
Av1MacroBlockD macroBlock,
Av1MacroBlockModeInfo modeInfo,
Point lumaOrigin,
Point chromaOrigin,
ushort tileIndex,
Av1PredictionMode lumaMode,
Av1TransformSize transformSize,
Av1TransformBlockContext blueContext,
Av1TransformBlockContext redContext,
Span<TSample> candidateBlueReconstruction,
Span<TSample> candidateRedReconstruction,
Span<int> candidateBlueCoefficients,
Span<int> candidateRedCoefficients,
Span<int> retainedBlueCoefficients,
Span<int> retainedRedCoefficients,
ref Av1EncoderTransformBlockState retainedBlueState,
ref Av1EncoderTransformBlockState retainedRedState,
ref long bestCost,
ref Av1EncoderPaletteInfo paletteInfo)
{
const Av1BlockSize BlockSize = Av1BlockSize.Block8x8;
const int LumaBlockLength = 8;
const int MaximumSampleCount = LumaBlockLength * LumaBlockLength;
ObuColorConfig colorConfig = this.picture.Sequence.SequenceHeader.ColorConfig;
int subsamplingX = colorConfig.SubSamplingX ? 1 : 0;
int subsamplingY = colorConfig.SubSamplingY ? 1 : 0;
int width = transformSize.GetWidth();
int height = transformSize.GetHeight();
ObuFrameSize frameSize = this.picture.Parent.FrameHeader.FrameSize;
int rows = Math.Min(LumaBlockLength, frameSize.FrameHeight - lumaOrigin.Y) >> subsamplingY;
int columns = Math.Min(LumaBlockLength, frameSize.FrameWidth - lumaOrigin.X) >> subsamplingX;
int activeSampleCount = rows * columns;
Span<short> blueSamples = stackalloc short[MaximumSampleCount];
Span<short> redSamples = stackalloc short[MaximumSampleCount];
blueSamples = blueSamples[..activeSampleCount];
redSamples = redSamples[..activeSampleCount];
Buffer2DRegion<TSample> blueSource = this.source.GetPlane(Av1Plane.U);
Buffer2DRegion<TSample> redSource = this.source.GetPlane(Av1Plane.V);
TOperator.CopyPaletteSamples(blueSource, chromaOrigin, rows, columns, blueSamples);
TOperator.CopyPaletteSamples(redSource, chromaOrigin, rows, columns, redSamples);
Span<short> uniqueBlueColors = stackalloc short[MaximumSampleCount];
Span<short> uniqueRedColors = stackalloc short[MaximumSampleCount];
int uniqueBlueColorCount = 0;
int uniqueRedColorCount = 0;
short blueMinimum = blueSamples[0];
short blueMaximum = blueSamples[0];
short redMinimum = redSamples[0];
short redMaximum = redSamples[0];
for (int sampleIndex = 0; sampleIndex < activeSampleCount; sampleIndex++)
{
short blueSample = blueSamples[sampleIndex];
short redSample = redSamples[sampleIndex];
if (!uniqueBlueColors[..uniqueBlueColorCount].Contains(blueSample))
{
uniqueBlueColors[uniqueBlueColorCount++] = blueSample;
}
if (!uniqueRedColors[..uniqueRedColorCount].Contains(redSample))
{
uniqueRedColors[uniqueRedColorCount++] = redSample;
}
blueMinimum = Math.Min(blueMinimum, blueSample);
blueMaximum = Math.Max(blueMaximum, blueSample);
redMinimum = Math.Min(redMinimum, redSample);
redMaximum = Math.Max(redMaximum, redSample);
}
int maximumColorCount = Math.Max(uniqueBlueColorCount, uniqueRedColorCount);
if (maximumColorCount < 2)
{
return false;
}
int maximumPaletteSize = Math.Min(maximumColorCount, Av1Constants.PaletteMaxSize);
Av1NeighborArrayUnit<Av1EncoderPaletteInfo> paletteContexts = this.picture.PaletteContexts[tileIndex];
Span<ushort> colorCache = stackalloc ushort[2 * Av1Constants.PaletteMaxSize];
int colorCacheSize = Av1TileWriter.GetPaletteCache(
paletteContexts,
macroBlock,
lumaOrigin,
Av1Plane.U,
colorCache);
colorCache = colorCache[..colorCacheSize];
int blockSizeContext = Av1TileWriter.GetPaletteBlockSizeContext(BlockSize);
bool hasLumaPalette = paletteInfo.PaletteSizes[0] != 0;
Buffer2DRegion<byte> colorIndexMap = this.superblock.Workspace
.GetPaletteMaps()
.GetMap(Av1PlaneType.Uv, width, height);
Span<byte> retainedColorIndexMap = stackalloc byte[MaximumSampleCount];
Span<short> blueCentroids = stackalloc short[Av1Constants.PaletteMaxSize];
Span<short> redCentroids = stackalloc short[Av1Constants.PaletteMaxSize];
Span<byte> colorIndices = stackalloc byte[MaximumSampleCount];
Span<TSample> bluePrediction = stackalloc TSample[MaximumSampleCount];
Span<TSample> redPrediction = stackalloc TSample[MaximumSampleCount];
Span<short> blueResidual = stackalloc short[MaximumSampleCount];
Span<short> redResidual = stackalloc short[MaximumSampleCount];
Span<ushort> bluePaletteColorStorage = stackalloc ushort[Av1Constants.PaletteMaxSize];
Span<ushort> redPaletteColorStorage = stackalloc ushort[Av1Constants.PaletteMaxSize];
int sampleCount = transformSize.GetSize2d();
int cacheThreshold = 4 << (this.bitDepth.GetBitCount() - 8);
bool paletteSelected = false;
// Chroma uses one paired K-means family; exhaustive size search avoids early header-cost pruning.
for (int paletteSize = 2; paletteSize <= maximumPaletteSize; paletteSize++)
{
Span<short> candidateBlueCentroids = blueCentroids[..paletteSize];
Span<short> candidateRedCentroids = redCentroids[..paletteSize];
Av1PaletteKMeans2D.InitializeCentroids(
blueMinimum,
blueMaximum,
redMinimum,
redMaximum,
candidateBlueCentroids,
candidateRedCentroids);
Av1PaletteKMeans2D.Cluster(
blueSamples,
redSamples,
candidateBlueCentroids,
candidateRedCentroids,
colorIndices[..activeSampleCount]);
for (int colorIndex = 0; colorIndex < paletteSize && !colorCache.IsEmpty; colorIndex++)
{
int minimumDifference = Math.Abs(candidateBlueCentroids[colorIndex] - colorCache[0]);
int nearestCacheIndex = 0;
for (int cacheIndex = 1; cacheIndex < colorCache.Length; cacheIndex++)
{
int difference = Math.Abs(candidateBlueCentroids[colorIndex] - colorCache[cacheIndex]);
if (difference < minimumDifference)
{
minimumDifference = difference;
nearestCacheIndex = cacheIndex;
}
}
if (minimumDifference <= cacheThreshold)
{
candidateBlueCentroids[colorIndex] = (short)colorCache[nearestCacheIndex];
}
}
// U is the coded ordering key, so each swap carries its paired V color with it.
for (int colorIndex = 0; colorIndex < paletteSize - 1; colorIndex++)
{
int minimumIndex = colorIndex;
for (int candidateIndex = colorIndex + 1; candidateIndex < paletteSize; candidateIndex++)
{
if (candidateBlueCentroids[candidateIndex] < candidateBlueCentroids[minimumIndex])
{
minimumIndex = candidateIndex;
}
}
if (minimumIndex != colorIndex)
{
(candidateBlueCentroids[colorIndex], candidateBlueCentroids[minimumIndex]) =
(candidateBlueCentroids[minimumIndex], candidateBlueCentroids[colorIndex]);
(candidateRedCentroids[colorIndex], candidateRedCentroids[minimumIndex]) =
(candidateRedCentroids[minimumIndex], candidateRedCentroids[colorIndex]);
}
}
Av1PaletteKMeans2D.AssignIndices(
blueSamples,
redSamples,
candidateBlueCentroids,
candidateRedCentroids,
colorIndices);
for (int row = 0; row < rows; row++)
{
Span<byte> mapRow = colorIndexMap.DangerousGetRowSpan(row)[..width];
colorIndices.Slice(row * columns, columns).CopyTo(mapRow);
mapRow[columns..].Fill(mapRow[columns - 1]);
}
// The shared U/V map covers the complete declared chroma block even at visible frame edges.
for (int row = rows; row < height; row++)
{
colorIndexMap.DangerousGetRowSpan(rows - 1)[..width]
.CopyTo(colorIndexMap.DangerousGetRowSpan(row));
}
Span<ushort> bluePaletteColors = bluePaletteColorStorage[..paletteSize];
Span<ushort> redPaletteColors = redPaletteColorStorage[..paletteSize];
for (int colorIndex = 0; colorIndex < paletteSize; colorIndex++)
{
bluePaletteColors[colorIndex] = (ushort)candidateBlueCentroids[colorIndex];
redPaletteColors[colorIndex] = (ushort)candidateRedCentroids[colorIndex];
}
TOperator.PreparePalette(
blueSource,
chromaOrigin,
bluePaletteColors,
colorIndexMap,
bluePrediction[..sampleCount],
blueResidual[..sampleCount],
transformSize);
TOperator.PreparePalette(
redSource,
chromaOrigin,
redPaletteColors,
colorIndexMap,
redPrediction[..sampleCount],
redResidual[..sampleCount],
transformSize);
Av1EncoderTransformBlockState candidateBlueState = default;
long distortion = TOperator.EncodePredictionCandidate(
this.blockWorkspace,
blueSource,
chromaOrigin,
bluePrediction,
blueResidual,
candidateBlueReconstruction,
candidateBlueCoefficients,
transformSize,
Av1TransformType.DctDct,
Av1Plane.U,
this.quantization.QIndex[0],
this.quantization.DeltaQDc[(int)Av1Plane.U],
this.quantization.DeltaQAc[(int)Av1Plane.U],
this.bitDepth,
ref candidateBlueState);
Av1EncoderTransformBlockState candidateRedState = default;
distortion += TOperator.EncodePredictionCandidate(
this.blockWorkspace,
redSource,
chromaOrigin,
redPrediction,
redResidual,
candidateRedReconstruction,
candidateRedCoefficients,
transformSize,
Av1TransformType.DctDct,
Av1Plane.V,
this.quantization.QIndex[0],
this.quantization.DeltaQDc[(int)Av1Plane.V],
this.quantization.DeltaQAc[(int)Av1Plane.V],
this.bitDepth,
ref candidateRedState);
int rate = Av1TileWriter.GetChromaModeCost(
writer,
this.picture.Parent.FrameHeader,
colorConfig,
modeInfo,
BlockSize,
lumaMode,
Av1ChromaPredictionMode.DC,
0);
rate += writer.GetPaletteUvModeCost(true, hasLumaPalette);
rate += writer.GetPaletteSizeCost(paletteSize, blockSizeContext, Av1PlaneType.Uv);
rate += Av1SymbolEncoder.GetPaletteUvColorCost(
colorCache,
bluePaletteColors,
redPaletteColors,
this.bitDepth.GetBitCount());
rate += writer.GetPaletteColorMapCost(
paletteSize,
Av1PlaneType.Uv,
rows,
columns,
colorIndexMap);
rate += writer.GetCoefficientCost(
transformSize,
Av1TransformType.DctDct,
lumaMode,
candidateBlueCoefficients,
Av1ComponentType.Chroma,
blueContext,
candidateBlueState.EndOfBlock,
this.picture.Parent.FrameHeader.UseReducedTransformSet,
Av1FilterIntraMode.AllFilterIntraModes);
rate += writer.GetCoefficientCost(
transformSize,
Av1TransformType.DctDct,
lumaMode,
candidateRedCoefficients,
Av1ComponentType.Chroma,
redContext,
candidateRedState.EndOfBlock,
this.picture.Parent.FrameHeader.UseReducedTransformSet,
Av1FilterIntraMode.AllFilterIntraModes);
long candidateCost = Av1RateDistortion.GetCost(this.rateMultiplier, rate, distortion);
if (candidateCost < bestCost)
{
Buffer2DRegion<TSample> blueReconstruction = this.reconstruction.GetPlane(Av1Plane.U);
Buffer2DRegion<TSample> redReconstruction = this.reconstruction.GetPlane(Av1Plane.V);
CopyCandidate(
candidateBlueReconstruction,
candidateBlueCoefficients,
blueReconstruction,
chromaOrigin,
retainedBlueCoefficients,
transformSize,
candidateBlueState,
ref retainedBlueState);
CopyCandidate(
candidateRedReconstruction,
candidateRedCoefficients,
redReconstruction,
chromaOrigin,
retainedRedCoefficients,
transformSize,
candidateRedState,
ref retainedRedState);
for (int row = 0; row < height; row++)
{
colorIndexMap.DangerousGetRowSpan(row)[..width]
.CopyTo(retainedColorIndexMap[(row * width)..]);
}
paletteInfo.PaletteSizes[1] = (byte)paletteSize;
paletteInfo.SetColors(Av1Plane.U, bluePaletteColors);
paletteInfo.SetColors(Av1Plane.V, redPaletteColors);
bestCost = candidateCost;
paletteSelected = true;
}
}
if (paletteSelected)
{
for (int row = 0; row < height; row++)
{
retainedColorIndexMap.Slice(row * width, width)
.CopyTo(colorIndexMap.DangerousGetRowSpan(row));
}
}
return paletteSelected;
}
}
}

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

@ -255,6 +255,7 @@ internal static partial class Av1IntraSuperblockEncoder
redCoefficients[this.codedAreaChroma..],
ref blueState,
ref redState,
ref paletteInfo,
out int chromaAngleDelta,
out byte chromaFromLumaIndex,
out sbyte chromaFromLumaSigns);

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

@ -780,6 +780,139 @@ public class Av1IntraSuperblockEncoderTests
initialSize: 256));
}
[Fact]
public void ProductionTileSelectsExactPairedChromaPalette()
{
AssertProductionTileSelectsExactPairedChromaPalette(useLumaPalette: false);
AssertProductionTileSelectsExactPairedChromaPalette(useLumaPalette: true);
}
private static void AssertProductionTileSelectsExactPairedChromaPalette(bool useLumaPalette)
{
const int Width = 8;
const int Height = 8;
const int QIndex = 37;
ObuColorConfig colorConfig = new()
{
IsMonochrome = false,
SubSamplingX = false,
SubSamplingY = false,
BitDepth = Av1BitDepth.EightBit
};
using Av1EncoderFrameBuffer<byte> source = new(
Configuration.Default,
Width,
Height,
8,
Av1ColorFormat.Yuv444,
0,
0);
using Av1EncoderFrameBuffer<byte> reconstruction = new(
Configuration.Default,
Width,
Height,
8,
Av1ColorFormat.Yuv444,
0,
0);
Buffer2DRegion<byte> lumaSource = source.Frame.CodedView.GetPlane(Av1Plane.Y);
Buffer2DRegion<byte> blueSource = source.Frame.CodedView.GetPlane(Av1Plane.U);
Buffer2DRegion<byte> redSource = source.Frame.CodedView.GetPlane(Av1Plane.V);
for (int row = 0; row < Height; row++)
{
Span<byte> lumaRow = lumaSource.DangerousGetRowSpan(row);
if (useLumaPalette)
{
for (int column = 0; column < Width; column++)
{
lumaRow[column] = column < Width / 2 ? (byte)64 : (byte)192;
}
}
else
{
lumaRow.Fill(128);
}
blueSource.DangerousGetRowSpan(row).Fill(row < Height / 2 ? (byte)32 : (byte)224);
redSource.DangerousGetRowSpan(row).Fill(row < Height / 2 ? (byte)200 : (byte)40);
}
ClearPlane(reconstruction.Luma);
ClearPlane(Assert.IsType<Buffer2D<byte>>(reconstruction.ChromaBlue));
ClearPlane(Assert.IsType<Buffer2D<byte>>(reconstruction.ChromaRed));
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 = new(
Configuration.Default,
source.Frame,
reconstruction.Frame,
picture.Picture,
coefficients,
superblockWorkspace,
blockWorkspace,
initialSize: 256);
ref Av1MacroBlockModeInfo mode = ref picture.Picture.GetMacroBlockModeInfo(default);
Assert.Equal(Av1ChromaPredictionMode.DC, mode.Block.UvMode);
Assert.Equal(useLumaPalette ? 2 : 0, superblockWorkspace.PaletteInfo.PaletteSizes[0]);
Assert.Equal(2, superblockWorkspace.PaletteInfo.PaletteSizes[1]);
Assert.Equal([32, 224], superblockWorkspace.PaletteInfo.GetColors(Av1Plane.U).ToArray());
Assert.Equal([200, 40], superblockWorkspace.PaletteInfo.GetColors(Av1Plane.V).ToArray());
Assert.Equal((ushort)0, coefficients.GetTransformBlockSpan(0, Av1Plane.U)[0].EndOfBlock);
Assert.Equal((ushort)0, coefficients.GetTransformBlockSpan(0, Av1Plane.V)[0].EndOfBlock);
Buffer2DRegion<byte> colorIndexMap = superblockWorkspace
.GetPaletteMaps()
.GetMap(Av1PlaneType.Uv, Width, Height);
Buffer2DRegion<byte> blueReconstruction = reconstruction.Frame.CodedView.GetPlane(Av1Plane.U);
Buffer2DRegion<byte> redReconstruction = reconstruction.Frame.CodedView.GetPlane(Av1Plane.V);
for (int row = 0; row < Height; row++)
{
byte expectedIndex = (byte)(row < Height / 2 ? 0 : 1);
foreach (byte index in colorIndexMap.DangerousGetRowSpan(row))
{
Assert.Equal(expectedIndex, index);
}
Assert.True(blueSource.DangerousGetRowSpan(row).SequenceEqual(blueReconstruction.DangerousGetRowSpan(row)));
Assert.True(redSource.DangerousGetRowSpan(row).SequenceEqual(redReconstruction.DangerousGetRowSpan(row)));
}
Assert.NotEqual(0, tileWriter.GetTileData(0).Length);
}
[Theory]
[InlineData((int)Av1PredictionMode.Vertical, 0)]
[InlineData((int)Av1PredictionMode.Horizontal, 0)]

Loading…
Cancel
Save