From bb335e46847519007a2374e0c4ecbf0728e3e64d Mon Sep 17 00:00:00 2001 From: James Jackson-South Date: Mon, 24 Aug 2026 18:28:49 +1000 Subject: [PATCH] Implement AV1 chroma-from-luma reconstruction --- HEIF_IMPLEMENTATION_PLAN.md | 2 +- .../Av1/Prediction/Av1PredictionDecoder.cs | 4 +- .../Av1ChromaFromLumaContext.cs | 172 +++++++++++++----- .../Heif/Av1/Tiling/Av1PartitionInfo.cs | 4 +- .../Formats/Heif/Av1/Tiling/Av1TileReader.cs | 2 +- .../Heif/Av1/Transform/Av1BlockDecoder.cs | 70 +++++-- .../Heif/Av1/Av1ChromaFromLumaTests.cs | 134 ++++++++++++++ 7 files changed, 315 insertions(+), 73 deletions(-) create mode 100644 tests/ImageSharp.Tests/Formats/Heif/Av1/Av1ChromaFromLumaTests.cs diff --git a/HEIF_IMPLEMENTATION_PLAN.md b/HEIF_IMPLEMENTATION_PLAN.md index 581bb522e..1af028959 100644 --- a/HEIF_IMPLEMENTATION_PLAN.md +++ b/HEIF_IMPLEMENTATION_PLAN.md @@ -58,7 +58,7 @@ This snapshot pins or classifies the available references and failures; it does | --- | --- | --- | --- | | `Av1YuvConverter.ConvertToRgb`, `ConvertFromRgb`, scalar row conversion, and chroma reconstruction | H.273 formulas 20-31 and the identity, YCgCo, and non-constant-luminance matrix formulas; AV1 section 6.4.2 chroma sample positions | libavif `src/reformat.c` and `src/colr.c` at `092276ce89098ead06db80975173191e5fee1826`; libaom `aom/aom_image.h` at `03087864cf4bea6abb0d28f95cf7843511413d8f` | Scalar behavioral oracle for 8-bit full/limited-range conversion. Decode covers monochrome, YUV 4:2:0, 4:2:2, and 4:4:4 with AV1 chroma sample positioning; encode remains YUV 4:4:4 at this snapshot. Later high-bit-depth and SIMD paths must match it. | | `Av1FrameBuffer` high-bit-depth sample layout and `Av1YuvConverter` 10/12-bit output conversion | AV1 section 6.4.1 bit depth and H.273 sample-range scaling | libaom `aom_scale/yv12config.h`, `av1/common/idct.c`, and `av1/common/reconintra.c` at `03087864cf4bea6abb0d28f95cf7843511413d8f`; libavif `src/avif.c` and `src/reformat.c` at `092276ce89098ead06db80975173191e5fee1826` | Establish two-byte native sample storage with sample-unit strides for 10/12-bit reconstruction and use the same scalar color model at every supported bit depth. | -| `Av1PredictionDecoder`, `Av1HighBitDepthPredictor`, and the scalar DC, directional, Paeth, smooth, and filter-intra predictors | AV1 sections 7.11.2 and 7.11.2.3 intra prediction | libaom `aom_dsp/intrapred.c` and `av1/common/reconintra.c` at `03087864cf4bea6abb0d28f95cf7843511413d8f` | Behavioral oracle for neighbor addressing, directional upsampling, Paeth selection, smooth normalization, filter-intra taps, high-bit-depth clipping, and chroma-from-luma row strides. Existing managed scalar tables and predictors remain the implementation base. The WIP rectangular byte-pipeline smooth digest expectations encode width/height-swapped weights and must be replaced only from an independently generated oracle, not regenerated from this implementation. | +| `Av1PredictionDecoder`, `Av1HighBitDepthPredictor`, `Av1ChromaFromLumaContext`, and the scalar DC, directional, Paeth, smooth, filter-intra, and chroma-from-luma predictors | AV1 sections 7.11.2 and 7.11.2.3 intra prediction | libaom `aom_dsp/intrapred.c`, `av1/common/reconintra.c`, `av1/common/cfl.c`, and `av1/common/cfl.h` at `03087864cf4bea6abb0d28f95cf7843511413d8f` | Behavioral oracle for neighbor addressing, directional upsampling, Paeth selection, smooth normalization, filter-intra taps, high-bit-depth clipping, chroma-from-luma storage/subsampling, and chroma-from-luma row strides. Existing managed scalar tables and predictors remain the implementation base. The WIP rectangular byte-pipeline smooth digest expectations encode width/height-swapped weights and must be replaced only from an independently generated oracle, not regenerated from this implementation. | | `Av1Inverse2dTransformer` and `Av1InverseTransformerFactory` | AV1 section 7.11.2 inverse transform and reconstruction | libaom `av1/common/av1_inv_txfm1d.c`, `av1/common/av1_inv_txfm2d.c`, and `av1/common/idct.c` at `03087864cf4bea6abb0d28f95cf7843511413d8f` | Scalar transform oracle for coefficient-row traversal, intermediate layout, stage ranges, clipping, and high-bit-depth sample addition. The managed 16-bit overload is also used as a parity oracle for the byte overload. | This table is intentionally incomplete. Add a row before each additional AV1 or HEVC algorithm is ported or materially reshaped. diff --git a/src/ImageSharp/Formats/Heif/Av1/Prediction/Av1PredictionDecoder.cs b/src/ImageSharp/Formats/Heif/Av1/Prediction/Av1PredictionDecoder.cs index e9269c7b2..d752e2009 100644 --- a/src/ImageSharp/Formats/Heif/Av1/Prediction/Av1PredictionDecoder.cs +++ b/src/ImageSharp/Formats/Heif/Av1/Prediction/Av1PredictionDecoder.cs @@ -159,7 +159,7 @@ internal class Av1PredictionDecoder // assert((transformSize.GetHeight() - 1) * CFL_BUF_LINE + transformSize.GetWidth() <= CFL_BUF_SQUARE); Av1BitDepth bitDepth = this.sequenceHeader.ColorConfig.BitDepth; ChromaFromLumaPredict( - chromaFromLumaContext.Q3Buffer!.DangerousGetSingleSpan(), + chromaFromLumaContext.Q3Buffer, pixelBuffer, stride, pixelBuffer, @@ -207,7 +207,7 @@ internal class Av1PredictionDecoder return Av1Math.RoundPowerOf2Signed(scaledLumaQ6, 6); } - private static void ChromaFromLumaPredict(Span predictedBufferQ3, Span predictedBuffer, int predictedStride, Span destinationBuffer, int destinationStride, int alphaQ3, Av1BitDepth bitDepth, int width, int height) + internal static void ChromaFromLumaPredict(Span predictedBufferQ3, Span predictedBuffer, int predictedStride, Span destinationBuffer, int destinationStride, int alphaQ3, Av1BitDepth bitDepth, int width, int height) where T : unmanaged, IBinaryInteger { // TODO: Make SIMD variant of this method. diff --git a/src/ImageSharp/Formats/Heif/Av1/Prediction/ChromaFromLuma/Av1ChromaFromLumaContext.cs b/src/ImageSharp/Formats/Heif/Av1/Prediction/ChromaFromLuma/Av1ChromaFromLumaContext.cs index c8112028d..9eb92971a 100644 --- a/src/ImageSharp/Formats/Heif/Av1/Prediction/ChromaFromLuma/Av1ChromaFromLumaContext.cs +++ b/src/ImageSharp/Formats/Heif/Av1/Prediction/ChromaFromLuma/Av1ChromaFromLumaContext.cs @@ -1,10 +1,9 @@ // Copyright (c) Six Labors. // Licensed under the Six Labors Split License. -using System.Runtime.CompilerServices; +using System.Numerics; using SixLabors.ImageSharp.Formats.Heif.Av1.OpenBitstreamUnit; using SixLabors.ImageSharp.Formats.Heif.Av1.Transform; -using SixLabors.ImageSharp.Memory; namespace SixLabors.ImageSharp.Formats.Heif.Av1.Prediction.ChromaFromLuma; @@ -17,62 +16,144 @@ internal class Av1ChromaFromLumaContext private readonly bool subX; private readonly bool subY; - public Av1ChromaFromLumaContext(Configuration configuration, ObuColorConfig colorConfig) + public Av1ChromaFromLumaContext(ObuColorConfig colorConfig) { this.subX = colorConfig.SubSamplingX; this.subY = colorConfig.SubSamplingY; - this.Q3Buffer = configuration.MemoryAllocator.Allocate2D(new Size(32, 32), AllocationOptions.Clean); + this.Q3Buffer = new short[BufferLine * BufferLine]; } - public Buffer2D Q3Buffer { get; private set; } + public short[] Q3Buffer { get; } public bool AreParametersComputed { get; private set; } + public void Store( + Span input, + int inputStride, + int row, + int column, + Av1TransformSize transformSize, + Av1BlockSize blockSize, + int modeInfoRow, + int modeInfoColumn) + where T : unmanaged, IBinaryInteger + { + if (blockSize.GetHeight() == 4 || blockSize.GetWidth() == 4) + { + // Subsampled chroma shares one CfL surface across the adjacent sub-8x8 luma blocks. + if ((modeInfoRow & 1) != 0 && this.subY) + { + row++; + } + + if ((modeInfoColumn & 1) != 0 && this.subX) + { + column++; + } + } + + int subX = this.subX ? 1 : 0; + int subY = this.subY ? 1 : 0; + int width = transformSize.GetWidth(); + int height = transformSize.GetHeight(); + int storeRow = row << (Av1Constants.ModeInfoSizeLog2 - subY); + int storeColumn = column << (Av1Constants.ModeInfoSizeLog2 - subX); + int storeWidth = width >> subX; + int storeHeight = height >> subY; + this.AreParametersComputed = false; + + if (column == 0 && row == 0) + { + this.bufferWidth = storeWidth; + this.bufferHeight = storeHeight; + } + else + { + this.bufferWidth = Math.Max(storeColumn + storeWidth, this.bufferWidth); + this.bufferHeight = Math.Max(storeRow + storeHeight, this.bufferHeight); + } + + int outputOffset = (storeRow * BufferLine) + storeColumn; + if (!this.subX) + { + // A direct luma sample is multiplied by eight to produce the Q3 representation used by CfL. + for (int y = 0; y < height; y++) + { + int inputRow = y * inputStride; + int outputRow = outputOffset + (y * BufferLine); + for (int x = 0; x < width; x++) + { + this.Q3Buffer[outputRow + x] = (short)(int.CreateChecked(input[inputRow + x]) << 3); + } + } + } + else if (!this.subY) + { + // The pair sum is multiplied by four, which is the Q3 representation of its horizontal average. + for (int y = 0; y < height; y++) + { + int inputRow = y * inputStride; + int outputRow = outputOffset + (y * BufferLine); + for (int x = 0; x < width; x += 2) + { + int sum = int.CreateChecked(input[inputRow + x]) + int.CreateChecked(input[inputRow + x + 1]); + this.Q3Buffer[outputRow + (x >> 1)] = (short)(sum << 2); + } + } + } + else + { + // The 2x2 sum is multiplied by two, which is the Q3 representation of its four-sample average. + for (int y = 0; y < height; y += 2) + { + int inputRow = y * inputStride; + int nextInputRow = inputRow + inputStride; + int outputRow = outputOffset + ((y >> 1) * BufferLine); + for (int x = 0; x < width; x += 2) + { + int sum = int.CreateChecked(input[inputRow + x]) + + int.CreateChecked(input[inputRow + x + 1]) + + int.CreateChecked(input[nextInputRow + x]) + + int.CreateChecked(input[nextInputRow + x + 1]); + + this.Q3Buffer[outputRow + (x >> 1)] = (short)(sum << 1); + } + } + } + } + public void ComputeParameters(Av1TransformSize transformSize) { Guard.IsFalse(this.AreParametersComputed, nameof(this.AreParametersComputed), "Do not call cfl_compute_parameters multiple time on the same values."); this.Pad(transformSize.GetWidth(), transformSize.GetHeight()); - SubtractAverage(ref this.Q3Buffer[0, 0], transformSize); + this.SubtractAverage(transformSize); this.AreParametersComputed = true; } private void Pad(int width, int height) { - int diff_width = width - this.bufferWidth; - int diff_height = height - this.bufferHeight; + int differenceWidth = width - this.bufferWidth; + int differenceHeight = height - this.bufferHeight; - if (diff_width > 0) + if (differenceWidth > 0) { - int min_height = height - diff_height; - ref short recon_buf_q3 = ref this.Q3Buffer[width - diff_width, 0]; - for (int j = 0; j < min_height; j++) + int minimumHeight = height - differenceHeight; + for (int y = 0; y < minimumHeight; y++) { - short last_pixel = Unsafe.Subtract(ref recon_buf_q3, 1); - Guard.IsTrue(Unsafe.IsAddressLessThan(ref Unsafe.Add(ref recon_buf_q3, diff_width), ref this.Q3Buffer[BufferLine, BufferLine]), nameof(recon_buf_q3), "Shall stay within bounds."); - for (int i = 0; i < diff_width; i++) - { - Unsafe.Add(ref recon_buf_q3, i) = last_pixel; - } - - recon_buf_q3 += BufferLine; + int rowOffset = y * BufferLine; + short lastPixel = this.Q3Buffer[rowOffset + this.bufferWidth - 1]; + this.Q3Buffer.AsSpan(rowOffset + this.bufferWidth, differenceWidth).Fill(lastPixel); } this.bufferWidth = width; } - if (diff_height > 0) + if (differenceHeight > 0) { - ref short recon_buf_q3 = ref this.Q3Buffer[0, height - diff_height]; - for (int j = 0; j < diff_height; j++) + for (int y = this.bufferHeight; y < height; y++) { - ref short last_row_q3 = ref Unsafe.Subtract(ref recon_buf_q3, BufferLine); - Guard.IsTrue(Unsafe.IsAddressLessThan(ref Unsafe.Add(ref recon_buf_q3, diff_width), ref this.Q3Buffer[BufferLine, BufferLine]), nameof(recon_buf_q3), "Shall stay within bounds."); - for (int i = 0; i < width; i++) - { - Unsafe.Add(ref recon_buf_q3, i) = Unsafe.Add(ref last_row_q3, i); - } - - recon_buf_q3 += BufferLine; + int rowOffset = y * BufferLine; + this.Q3Buffer.AsSpan(rowOffset - BufferLine, width).CopyTo(this.Q3Buffer.AsSpan(rowOffset, width)); } this.bufferHeight = height; @@ -83,38 +164,31 @@ internal class Av1ChromaFromLumaContext * svt_subtract_average_c * Calculate the DC value by averaging over all sample. Subtract DC value to get AC values In C ************************************************************************************************/ - private static void SubtractAverage(ref short pred_buf_q3, Av1TransformSize transformSize) + private void SubtractAverage(Av1TransformSize transformSize) { int width = transformSize.GetWidth(); int height = transformSize.GetHeight(); int roundOffset = (width * height) >> 1; int pelCountLog2 = transformSize.GetBlockWidthLog2() + transformSize.GetBlockHeightLog2(); - int sum_q3 = 0; - ref short pred_buf = ref pred_buf_q3; - for (int j = 0; j < height; j++) + int sumQ3 = roundOffset; + for (int y = 0; y < height; y++) { - // assert(pred_buf_q3 + tx_width <= cfl->pred_buf_q3 + CFL_BUF_SQUARE); - for (int i = 0; i < width; i++) + int rowOffset = y * BufferLine; + for (int x = 0; x < width; x++) { - sum_q3 += Unsafe.Add(ref pred_buf, i); + sumQ3 += this.Q3Buffer[rowOffset + x]; } - - pred_buf += BufferLine; } - int avg_q3 = (sum_q3 + roundOffset) >> pelCountLog2; + int averageQ3 = sumQ3 >> pelCountLog2; - // Loss is never more than 1/2 (in Q3) - // assert(abs((avg_q3 * (1 << num_pel_log2)) - sum_q3) <= 1 << num_pel_log2 >> - // 1); - for (int j = 0; j < height; j++) + for (int y = 0; y < height; y++) { - for (int i = 0; i < width; i++) + int rowOffset = y * BufferLine; + for (int x = 0; x < width; x++) { - Unsafe.Add(ref pred_buf_q3, i) -= (short)avg_q3; + this.Q3Buffer[rowOffset + x] -= (short)averageQ3; } - - pred_buf_q3 += BufferLine; } } } diff --git a/src/ImageSharp/Formats/Heif/Av1/Tiling/Av1PartitionInfo.cs b/src/ImageSharp/Formats/Heif/Av1/Tiling/Av1PartitionInfo.cs index 470f0cd53..159a06e41 100644 --- a/src/ImageSharp/Formats/Heif/Av1/Tiling/Av1PartitionInfo.cs +++ b/src/ImageSharp/Formats/Heif/Av1/Tiling/Av1PartitionInfo.cs @@ -87,7 +87,7 @@ internal class Av1PartitionInfo public Av1ChromaFromLumaContext? ChromaFromLumaContext { get; internal set; } - public void ComputeBoundaryOffsets(Configuration configuration, ObuSequenceHeader sequenceHeader, ObuFrameHeader frameHeader, Av1TileInfo tileInfo) + public void ComputeBoundaryOffsets(ObuSequenceHeader sequenceHeader, ObuFrameHeader frameHeader, Av1TileInfo tileInfo) { Av1BlockSize blockSize = this.ModeInfo.BlockSize; int bw4 = blockSize.Get4x4WideCount(); @@ -118,8 +118,6 @@ internal class Av1PartitionInfo // For V plane chroma bock this.WidthInPixels[2] = Math.Max(1, bw4 >> subX) * modeInfoSize; this.HeightInPixels[2] = Math.Max(1, bh4 >> subY) * modeInfoSize; - - this.ChromaFromLumaContext = new Av1ChromaFromLumaContext(configuration, sequenceHeader.ColorConfig); } public int GetMaxBlockWide(Av1BlockSize blockSize, bool subX) diff --git a/src/ImageSharp/Formats/Heif/Av1/Tiling/Av1TileReader.cs b/src/ImageSharp/Formats/Heif/Av1/Tiling/Av1TileReader.cs index b0395736f..f8a735ffd 100644 --- a/src/ImageSharp/Formats/Heif/Av1/Tiling/Av1TileReader.cs +++ b/src/ImageSharp/Formats/Heif/Av1/Tiling/Av1TileReader.cs @@ -303,7 +303,7 @@ internal class Av1TileReader : IAv1TileReader partitionInfo.ColumnIndex = columnIndex; partitionInfo.RowIndex = rowIndex; superblockInfo.BlockCount++; - partitionInfo.ComputeBoundaryOffsets(this.configuration, this.SequenceHeader, this.FrameHeader, tileInfo); + partitionInfo.ComputeBoundaryOffsets(this.SequenceHeader, this.FrameHeader, tileInfo); if (hasChroma) { if (this.SequenceHeader.ColorConfig.SubSamplingY && block4x4Height == 1) diff --git a/src/ImageSharp/Formats/Heif/Av1/Transform/Av1BlockDecoder.cs b/src/ImageSharp/Formats/Heif/Av1/Transform/Av1BlockDecoder.cs index 54bb7f4b2..96028ab96 100644 --- a/src/ImageSharp/Formats/Heif/Av1/Transform/Av1BlockDecoder.cs +++ b/src/ImageSharp/Formats/Heif/Av1/Transform/Av1BlockDecoder.cs @@ -5,6 +5,7 @@ using System.Runtime.CompilerServices; using SixLabors.ImageSharp.Formats.Heif.Av1.OpenBitstreamUnit; using SixLabors.ImageSharp.Formats.Heif.Av1.Pipeline.Quantification; using SixLabors.ImageSharp.Formats.Heif.Av1.Prediction; +using SixLabors.ImageSharp.Formats.Heif.Av1.Prediction.ChromaFromLuma; using SixLabors.ImageSharp.Formats.Heif.Av1.Tiling; namespace SixLabors.ImageSharp.Formats.Heif.Av1.Transform; @@ -23,6 +24,8 @@ internal class Av1BlockDecoder private readonly int[] currentCoefficientIndex; + private readonly Av1ChromaFromLumaContext chromaFromLumaContext; + public Av1BlockDecoder(ObuSequenceHeader sequenceHeader, ObuFrameHeader frameHeader, Av1FrameInfo frameInfo, Av1FrameBuffer frameBuffer) { this.sequenceHeader = sequenceHeader; @@ -36,6 +39,7 @@ internal class Av1BlockDecoder this.CurrentInverseQuantizationCoefficients = new int[inverseQuantizationSize]; this.isLoopFilterEnabled = false; this.currentCoefficientIndex = new int[3]; + this.chromaFromLumaContext = new(sequenceHeader.ColorConfig); } public int[] CurrentInverseQuantizationCoefficients { get; private set; } @@ -57,7 +61,27 @@ internal class Av1BlockDecoder Av1TransformSize transformSize; int transformUnitCount; bool hasChroma = Av1TileReader.HasChroma(this.sequenceHeader, modeInfoPosition, blockSize); - Av1PartitionInfo partitionInfo = new(modeInfo, superblockInfo, hasChroma, Av1PartitionType.None); + Av1PartitionInfo partitionInfo = new(modeInfo, superblockInfo, hasChroma, modeInfo.PartitionType) + { + ColumnIndex = modeInfoPosition.X, + RowIndex = modeInfoPosition.Y, + ChromaFromLumaContext = this.chromaFromLumaContext + }; + + partitionInfo.ComputeBoundaryOffsets(this.sequenceHeader, this.frameHeader, tileInfo); + + if (hasChroma) + { + if (colorConfig.SubSamplingY && blockSize.Get4x4HighCount() == 1) + { + partitionInfo.AvailableAboveForChroma = modeInfoPosition.Y - 2 >= tileInfo.ModeInfoRowStart; + } + + if (colorConfig.SubSamplingX && blockSize.Get4x4WideCount() == 1) + { + partitionInfo.AvailableLeftForChroma = modeInfoPosition.X - 2 >= tileInfo.ModeInfoColumnStart; + } + } int maxBlocksWide = partitionInfo.GetMaxBlockWide(blockSize, false); int maxBlocksHigh = partitionInfo.GetMaxBlockHigh(blockSize, false); @@ -234,22 +258,32 @@ internal class Av1BlockDecoder } // Store Luma for CFL if required! - if (plane == (int)Av1Plane.Y && StoreChromeFromLumeRequired(colorConfig, partitionInfo, hasChroma)) + if (plane == (int)Av1Plane.Y && StoreChromaFromLumaRequired(colorConfig, partitionInfo)) { - /* - // SVT: svt_cfl_store_tx - ChromaFromLumaStoreTransform( - partitionInfo, - this.chromaFromLumaContext, - transformInfo.OffsetY, - transformInfo.OffsetX, - transformSize, - blockSize, - colorConfig, - transformBlockReconstructionBuffer, - reconstructionStride, - is16BitsPipeline); - */ + if (highBitDepth) + { + this.chromaFromLumaContext.Store( + highBitDepthTransformBlockReconstructionBuffer[reconstructionStride..], + reconstructionStride, + transformInfo[0].OffsetY, + transformInfo[0].OffsetX, + transformSize, + blockSize, + modeInfoPosition.Y, + modeInfoPosition.X); + } + else + { + this.chromaFromLumaContext.Store( + transformBlockReconstructionBuffer[reconstructionStride..], + reconstructionStride, + transformInfo[0].OffsetY, + transformInfo[0].OffsetX, + transformSize, + blockSize, + modeInfoPosition.Y, + modeInfoPosition.X); + } } // increment transform pointer @@ -319,5 +353,7 @@ internal class Av1BlockDecoder } } - private static bool StoreChromeFromLumeRequired(ObuColorConfig colorConfig, Av1PartitionInfo partitionInfo, bool hasChroma) => false; + private static bool StoreChromaFromLumaRequired(ObuColorConfig colorConfig, Av1PartitionInfo partitionInfo) + => !colorConfig.IsMonochrome && + (!partitionInfo.IsChroma || partitionInfo.ModeInfo.UvMode == Av1PredictionMode.UvChromaFromLuma); } diff --git a/tests/ImageSharp.Tests/Formats/Heif/Av1/Av1ChromaFromLumaTests.cs b/tests/ImageSharp.Tests/Formats/Heif/Av1/Av1ChromaFromLumaTests.cs new file mode 100644 index 000000000..ca172fb21 --- /dev/null +++ b/tests/ImageSharp.Tests/Formats/Heif/Av1/Av1ChromaFromLumaTests.cs @@ -0,0 +1,134 @@ +// Copyright (c) Six Labors. +// Licensed under the Six Labors Split License. + +using SixLabors.ImageSharp.Formats.Heif.Av1; +using SixLabors.ImageSharp.Formats.Heif.Av1.OpenBitstreamUnit; +using SixLabors.ImageSharp.Formats.Heif.Av1.Prediction; +using SixLabors.ImageSharp.Formats.Heif.Av1.Prediction.ChromaFromLuma; +using SixLabors.ImageSharp.Formats.Heif.Av1.Tiling; +using SixLabors.ImageSharp.Formats.Heif.Av1.Transform; + +namespace SixLabors.ImageSharp.Tests.Formats.Heif.Av1; + +[Trait("Format", "Avif")] +public class Av1ChromaFromLumaTests +{ + [Theory] + [InlineData(false, false, new short[] { 8, 16, 24, 32, 40, 48, 56, 64, 72, 80, 88, 96, 104, 112, 120, 128 })] + [InlineData(true, false, new short[] { 12, 28, 44, 60, 76, 92, 108, 124 })] + [InlineData(true, true, new short[] { 28, 44, 92, 108 })] + public void Store8BitMatchesLibaomSubsampling(bool subX, bool subY, short[] expected) + { + ObuColorConfig colorConfig = new() { SubSamplingX = subX, SubSamplingY = subY }; + Av1ChromaFromLumaContext context = new(colorConfig); + byte[] input = Enumerable.Range(1, 16).Select(x => (byte)x).ToArray(); + + context.Store(input, 4, 0, 0, Av1TransformSize.Size4x4, Av1BlockSize.Block4x4, 0, 0); + + int width = 4 >> (subX ? 1 : 0); + int height = 4 >> (subY ? 1 : 0); + Assert.Equal(expected, GetBlock(context.Q3Buffer, width, height)); + } + + [Fact] + public void StoreHighBitDepthPreservesTwelveBitQ3Range() + { + ObuColorConfig colorConfig = new(); + Av1ChromaFromLumaContext context = new(colorConfig); + short[] input = Enumerable.Repeat((short)4095, 16).ToArray(); + + context.Store(input, 4, 0, 0, Av1TransformSize.Size4x4, Av1BlockSize.Block4x4, 0, 0); + + Assert.All(GetBlock(context.Q3Buffer, 4, 4), value => Assert.Equal(32760, value)); + } + + [Fact] + public void StoreCombinesSub8x8LumaBeforeSubtractingAverage() + { + ObuColorConfig colorConfig = new() { SubSamplingX = true, SubSamplingY = true }; + Av1ChromaFromLumaContext context = new(colorConfig); + + context.Store(Enumerable.Repeat((byte)10, 16).ToArray(), 4, 0, 0, Av1TransformSize.Size4x4, Av1BlockSize.Block4x4, 0, 0); + context.Store(Enumerable.Repeat((byte)20, 16).ToArray(), 4, 0, 0, Av1TransformSize.Size4x4, Av1BlockSize.Block4x4, 0, 1); + context.Store(Enumerable.Repeat((byte)30, 16).ToArray(), 4, 0, 0, Av1TransformSize.Size4x4, Av1BlockSize.Block4x4, 1, 0); + context.Store(Enumerable.Repeat((byte)40, 16).ToArray(), 4, 0, 0, Av1TransformSize.Size4x4, Av1BlockSize.Block4x4, 1, 1); + + context.ComputeParameters(Av1TransformSize.Size4x4); + + short[] expected = + [ + -120, -120, -40, -40, + -120, -120, -40, -40, + 40, 40, 120, 120, + 40, 40, 120, 120 + ]; + + Assert.Equal(expected, GetBlock(context.Q3Buffer, 4, 4)); + } + + [Fact] + public void ComputeParametersPadsFrameEdgeBeforeSubtractingAverage() + { + ObuColorConfig colorConfig = new(); + Av1ChromaFromLumaContext context = new(colorConfig); + byte[] input = Enumerable.Range(1, 16).Select(x => (byte)x).ToArray(); + context.Store(input, 4, 0, 0, Av1TransformSize.Size4x4, Av1BlockSize.Block4x4, 0, 0); + + context.ComputeParameters(Av1TransformSize.Size8x8); + + short[] actual = GetBlock(context.Q3Buffer, 8, 8); + Assert.Equal(-90, actual[0]); + Assert.Equal(-66, actual[7]); + Assert.Equal(6, actual[56]); + Assert.Equal(30, actual[63]); + Assert.Equal(0, actual.Sum(x => x)); + } + + [Fact] + public void Predict8BitAddsScaledLumaAndClips() + { + short[] lumaQ3 = new short[32 * 32]; + new short[] { -64, -32, 64, 64 }.CopyTo(lumaQ3, 0); + byte[] dcPrediction = [0, 128, 250, 255]; + byte[] destination = new byte[4]; + + Av1PredictionDecoder.ChromaFromLumaPredict(lumaQ3, dcPrediction, 4, destination, 4, 8, Av1BitDepth.EightBit, 4, 1); + + Assert.Equal(new byte[] { 0, 124, 255, 255 }, destination); + } + + [Theory] + [InlineData((int)Av1BitDepth.TenBit, 1023)] + [InlineData((int)Av1BitDepth.TwelveBit, 4095)] + public void PredictHighBitDepthAddsScaledLumaAndClips(int bitDepthIndex, short maximum) + { + short[] lumaQ3 = new short[32 * 32]; + new short[] { -128, -64, 64, 128 }.CopyTo(lumaQ3, 0); + short[] dcPrediction = [5, (short)(maximum / 2), (short)(maximum - 5), maximum]; + short[] destination = new short[4]; + + Av1PredictionDecoder.ChromaFromLumaPredict( + lumaQ3, + dcPrediction, + 4, + destination, + 4, + 16, + (Av1BitDepth)bitDepthIndex, + 4, + 1); + + Assert.Equal(new short[] { 0, (short)((maximum / 2) - 16), maximum, maximum }, destination); + } + + private static short[] GetBlock(short[] buffer, int width, int height) + { + short[] result = new short[width * height]; + for (int y = 0; y < height; y++) + { + buffer.AsSpan(y * 32, width).CopyTo(result.AsSpan(y * width, width)); + } + + return result; + } +}