Test case for coefficient round trip

SixLabors · Nov 23, 2024 · fa886a4 · fa886a4
1 parent 5aca407
commit fa886a4
Show file tree

Hide file tree

Showing 12 changed files with 303 additions and 233 deletions.
diff --git a/src/ImageSharp/Formats/Heif/Av1/Entropy/Av1SymbolDecoder.cs b/src/ImageSharp/Formats/Heif/Av1/Entropy/Av1SymbolDecoder.cs
@@ -206,6 +206,88 @@ public int ReadChromaFromLumaAlphaV(int jointSignPlus1)
         return r.ReadSymbol(this.chromeForLumaAlpha[context]);
     }
 
+    /// <summary>
+    /// 5.11.39. Coefficients syntax.
+    /// </summary>
+    /// <remarks>
+    /// The implementation is taken from SVT-AV1 library, which deviates from the code flow in the specification.
+    /// </remarks>
+    public int ReadCoefficients(
+        Av1BlockModeInfo modeInfo,
+        Point blockPosition,
+        int[] aboveContexts,
+        int[] leftContexts,
+        int aboveOffset,
+        int leftOffset,
+        int plane,
+        int blocksWide,
+        int blocksHigh,
+        Av1TransformBlockContext transformBlockContext,
+        Av1TransformSize transformSize,
+        bool isLossless,
+        bool useReducedTransformSet,
+        Av1TransformInfo transformInfo,
+        int modeBlocksToRightEdge,
+        int modeBlocksToBottomEdge,
+        Span<int> coefficientBuffer)
+    {
+        int width = transformSize.GetWidth();
+        int height = transformSize.GetHeight();
+        Av1TransformSize transformSizeContext = (Av1TransformSize)(((int)transformSize.GetSquareSize() + ((int)transformSize.GetSquareUpSize() + 1)) >> 1);
+        Av1PlaneType planeType = (Av1PlaneType)Math.Min(plane, 1);
+        int culLevel = 0;
+
+        byte[] levelsBuffer = new byte[Av1Constants.TransformPad2d];
+        Span<byte> levels = levelsBuffer.AsSpan()[(Av1Constants.TransformPadTop * (width + Av1Constants.TransformPadHorizontal))..];
+
+        bool allZero = this.ReadTransformBlockSkip(transformSizeContext, transformBlockContext.SkipContext);
+        int bwl = transformSize.GetBlockWidthLog2();
+        int endOfBlock;
+        if (allZero)
+        {
+            if (plane == 0)
+            {
+                transformInfo.Type = Av1TransformType.DctDct;
+                transformInfo.CodeBlockFlag = false;
+            }
+
+            this.UpdateCoefficientContext(modeInfo, aboveContexts, leftContexts, blocksWide, blocksHigh, transformSize, blockPosition, aboveOffset, leftOffset, culLevel, modeBlocksToRightEdge, modeBlocksToBottomEdge);
+            return 0;
+        }
+
+        transformInfo.Type = ComputeTransformType(planeType, modeInfo, isLossless, transformSize, transformInfo, useReducedTransformSet);
+        Av1TransformClass transformClass = transformInfo.Type.ToClass();
+        Av1ScanOrder scanOrder = Av1ScanOrderConstants.GetScanOrder(transformSize, transformInfo.Type);
+        ReadOnlySpan<short> scan = scanOrder.Scan;
+
+        endOfBlock = this.ReadEndOfBlockPosition(transformSize, transformClass, transformSizeContext, planeType);
+        if (endOfBlock > 1)
+        {
+            Array.Fill(levelsBuffer, (byte)0, 0, ((width + Av1Constants.TransformPadHorizontal) * (height + Av1Constants.TransformPadVertical)) + Av1Constants.TransformPadEnd);
+        }
+
+        this.ReadCoefficientsEndOfBlock(transformClass, endOfBlock, height, scan, bwl, levels, transformSizeContext, planeType);
+        if (endOfBlock > 1)
+        {
+            if (transformClass == Av1TransformClass.Class2D)
+            {
+                this.ReadCoefficientsReverse2d(transformSize, 1, endOfBlock - 1 - 1, scan, bwl, levels, transformSizeContext, planeType);
+                this.ReadCoefficientsReverse(transformSize, transformClass, 0, 0, scan, bwl, levels, transformSizeContext, planeType);
+            }
+            else
+            {
+                this.ReadCoefficientsReverse(transformSize, transformClass, 0, endOfBlock - 1 - 1, scan, bwl, levels, transformSizeContext, planeType);
+            }
+        }
+
+        DebugGuard.MustBeGreaterThan(scan.Length, 0, nameof(scan));
+        culLevel = this.ReadCoefficientsDc(coefficientBuffer, endOfBlock, scan, bwl, levels, transformBlockContext.DcSignContext, planeType);
+        this.UpdateCoefficientContext(modeInfo, aboveContexts, leftContexts, blocksWide, blocksHigh, transformSize, blockPosition, aboveOffset, leftOffset, culLevel, modeBlocksToRightEdge, modeBlocksToBottomEdge);
+
+        transformInfo.CodeBlockFlag = true;
+        return endOfBlock;
+    }
+
     public int ReadEndOfBlockPosition(Av1TransformSize transformSize, Av1TransformClass transformClass, Av1TransformSize transformSizeContext, Av1PlaneType planeType)
     {
         int endOfBlockExtra = 0;
@@ -581,6 +663,106 @@ private static void SetDcSign(ref int culLevel, int dcValue)
         }
     }
 
+
+    private void UpdateCoefficientContext(
+        Av1BlockModeInfo modeInfo,
+        int[] aboveContexts,
+        int[] leftContexts,
+        int blocksWide,
+        int blocksHigh,
+        Av1TransformSize transformSize,
+        Point blockPosition,
+        int aboveOffset,
+        int leftOffset,
+        int culLevel,
+        int modeBlockToRightEdge,
+        int modeBlockToBottomEdge)
+    {
+        int transformSizeWide = transformSize.Get4x4WideCount();
+        int transformSizeHigh = transformSize.Get4x4HighCount();
+
+        if (modeBlockToRightEdge < 0)
+        {
+            int aboveContextCount = Math.Min(transformSizeWide, blocksWide - aboveOffset);
+            Array.Fill(aboveContexts, culLevel, 0, aboveContextCount);
+            Array.Fill(aboveContexts, 0, aboveContextCount, transformSizeWide - aboveContextCount);
+        }
+        else
+        {
+            Array.Fill(aboveContexts, culLevel, 0, transformSizeWide);
+        }
+
+        if (modeBlockToBottomEdge < 0)
+        {
+            int leftContextCount = Math.Min(transformSizeHigh, blocksHigh - leftOffset);
+            Array.Fill(leftContexts, culLevel, 0, leftContextCount);
+            Array.Fill(leftContexts, 0, leftContextCount, transformSizeWide - leftContextCount);
+        }
+        else
+        {
+            Array.Fill(leftContexts, culLevel, 0, transformSizeHigh);
+        }
+    }
+
+    private static Av1TransformType ComputeTransformType(Av1PlaneType planeType, Av1BlockModeInfo modeInfo, bool isLossless, Av1TransformSize transformSize, Av1TransformInfo transformInfo, bool useReducedTransformSet)
+    {
+        Av1TransformType transformType = Av1TransformType.DctDct;
+        if (isLossless || transformSize.GetSquareUpSize() > Av1TransformSize.Size32x32)
+        {
+            transformType = Av1TransformType.DctDct;
+        }
+        else
+        {
+            if (planeType == Av1PlaneType.Y)
+            {
+                transformType = transformInfo.Type;
+            }
+            else
+            {
+                // In intra mode, uv planes don't share the same prediction mode as y
+                // plane, so the tx_type should not be shared
+                transformType = ConvertIntraModeToTransformType(modeInfo, Av1PlaneType.Uv);
+            }
+        }
+
+        Av1TransformSetType transformSetType = GetExtendedTransformSetType(transformSize, useReducedTransformSet);
+        if (!transformType.IsExtendedSetUsed(transformSetType))
+        {
+            transformType = Av1TransformType.DctDct;
+        }
+
+        return transformType;
+    }
+
+    private static Av1TransformSetType GetExtendedTransformSetType(Av1TransformSize transformSize, bool useReducedSet)
+    {
+        Av1TransformSize squareUpSize = transformSize.GetSquareUpSize();
+
+        if (squareUpSize >= Av1TransformSize.Size32x32)
+        {
+            return Av1TransformSetType.DctOnly;
+        }
+
+        if (useReducedSet)
+        {
+            return Av1TransformSetType.Dtt4Identity;
+        }
+
+        Av1TransformSize squareSize = transformSize.GetSquareSize();
+        return squareSize == Av1TransformSize.Size16x16 ? Av1TransformSetType.Dtt4Identity : Av1TransformSetType.Dtt4Identity1dDct;
+    }
+
+    private static Av1TransformType ConvertIntraModeToTransformType(Av1BlockModeInfo modeInfo, Av1PlaneType planeType)
+    {
+        Av1PredictionMode mode = (planeType == Av1PlaneType.Y) ? modeInfo.YMode : modeInfo.UvMode;
+        if (mode == Av1PredictionMode.UvChromaFromLuma)
+        {
+            mode = Av1PredictionMode.DC;
+        }
+
+        return mode.ToTransformType();
+    }
+
     internal static Av1Distribution GetSplitOrHorizontalDistribution(Av1Distribution[] inputs, Av1BlockSize blockSize, int context)
     {
         Av1Distribution input = inputs[context];

diff --git a/src/ImageSharp/Formats/Heif/Av1/Entropy/Av1SymbolEncoder.cs b/src/ImageSharp/Formats/Heif/Av1/Entropy/Av1SymbolEncoder.cs
@@ -415,13 +415,13 @@ private static void InitializeLevels(Span<int> coefficientBuffer, int width, int
         ref byte ls = ref levels[0];
 
         Unsafe.InitBlock(ref levels[-Av1Constants.TransformPadTop * stride], 0, (uint)(Av1Constants.TransformPadTop * stride));
-        Unsafe.InitBlock(ref levels[stride * height], 0, (uint)(Av1Constants.TransformPadBottom * stride + Av1Constants.TransformPadEnd));
+        Unsafe.InitBlock(ref levels[stride * height], 0, (uint)((Av1Constants.TransformPadBottom * stride) + Av1Constants.TransformPadEnd));
 
-        for (int i = 0; i < height; i++)
+        for (int y = 0; y < height; y++)
         {
-            for (int j = 0; j < width; j++)
+            for (int x = 0; x < width; x++)
             {
-                ls = (byte)Av1Math.Clamp(Math.Abs(coefficientBuffer[i * width + j]), 0, byte.MaxValue);
+                ls = (byte)Av1Math.Clamp(Math.Abs(coefficientBuffer[(y * width) + x]), 0, byte.MaxValue);
                 ls = ref Unsafe.Add(ref ls, 1);
             }
 
@@ -433,7 +433,10 @@ private static void InitializeLevels(Span<int> coefficientBuffer, int width, int
     /// SVT: set_levels from EbCommonUtils.h
     /// </summary>
     private static Span<byte> SetLevels(Span<byte> levelsBuffer, int width)
-        => levelsBuffer.Slice(Av1Constants.TransformPadTop * (width + Av1Constants.TransformPadHorizontal));
+    {
+        int stride = width + Av1Constants.TransformPadHorizontal;
+        return levelsBuffer[(Av1Constants.TransformPadTop * stride)..];
+    }
 
     private void WriteSkip(bool hasEndOfBlock, int context)
     {

diff --git a/src/ImageSharp/Formats/Heif/Av1/Tiling/Av1BlockModeInfoEncoder.cs b/src/ImageSharp/Formats/Heif/Av1/Tiling/Av1BlockModeInfoEncoder.cs
@@ -1,30 +1,27 @@
-// Copyright (c) Six Labors.
+// Copyright (c) Six Labors.
 // Licensed under the Six Labors Split License.
 
 using SixLabors.ImageSharp.Formats.Heif.Av1.Prediction;
 
 namespace SixLabors.ImageSharp.Formats.Heif.Av1.Tiling;
 
-internal partial class Av1TileWriter
+internal class Av1BlockModeInfoEncoder
 {
-    internal class Av1BlockModeInfoEncoder
-    {
-        public Av1BlockSize BlockSize { get; }
+    public Av1BlockSize BlockSize { get; }
 
-        public Av1PredictionMode PredictionMode { get; }
+    public Av1PredictionMode PredictionMode { get; }
 
-        public Av1PartitionType PartitionType { get; }
+    public Av1PartitionType PartitionType { get; }
 
-        public Av1PredictionMode UvPredictionMode { get; }
+    public Av1PredictionMode UvPredictionMode { get; }
 
-        public bool Skip { get; } = true;
+    public bool Skip { get; } = true;
 
-        public bool SkipMode { get; } = true;
+    public bool SkipMode { get; } = true;
 
-        public bool UseIntraBlockCopy { get; } = true;
+    public bool UseIntraBlockCopy { get; } = true;
 
-        public int SegmentId { get; }
+    public int SegmentId { get; }
 
-        public int TransformDepth { get; internal set; }
-    }
+    public int TransformDepth { get; internal set; }
 }
diff --git a/src/ImageSharp/Formats/Heif/Av1/Tiling/Av1Common.cs b/src/ImageSharp/Formats/Heif/Av1/Tiling/Av1Common.cs
@@ -5,18 +5,15 @@
 
 namespace SixLabors.ImageSharp.Formats.Heif.Av1.Tiling;
 
-internal partial class Av1TileWriter
+internal class Av1Common
 {
-    internal class Av1Common
-    {
-        public int ModeInfoRowCount { get; internal set; }
+    public int ModeInfoRowCount { get; internal set; }
 
-        public int ModeInfoColumnCount { get; internal set; }
+    public int ModeInfoColumnCount { get; internal set; }
 
-        public int ModeInfoStride { get; internal set; }
+    public int ModeInfoStride { get; internal set; }
 
-        public required ObuFrameSize FrameSize { get; internal set; }
+    public required ObuFrameSize FrameSize { get; internal set; }
 
-        public required ObuTileGroupHeader TilesInfo { get; internal set; }
-    }
+    public required ObuTileGroupHeader TilesInfo { get; internal set; }
 }
diff --git a/src/ImageSharp/Formats/Heif/Av1/Tiling/Av1MacroBlockModeInfo.cs b/src/ImageSharp/Formats/Heif/Av1/Tiling/Av1MacroBlockModeInfo.cs
@@ -3,14 +3,11 @@
 
 namespace SixLabors.ImageSharp.Formats.Heif.Av1.Tiling;
 
-internal partial class Av1TileWriter
+internal class Av1MacroBlockModeInfo
 {
-    internal class Av1MacroBlockModeInfo
-    {
-        public required Av1BlockModeInfoEncoder Block { get; internal set; }
+    public required Av1BlockModeInfoEncoder Block { get; internal set; }
 
-        public required Av1PaletteLumaModeInfo Palette { get; internal set; }
+    public required Av1PaletteLumaModeInfo Palette { get; internal set; }
 
-        public int CdefStrength { get; internal set; }
-    }
+    public int CdefStrength { get; internal set; }
 }
diff --git a/src/ImageSharp/Formats/Heif/Av1/Tiling/Av1NeighborArrayUnit.cs b/src/ImageSharp/Formats/Heif/Av1/Tiling/Av1NeighborArrayUnit.cs
@@ -11,14 +11,6 @@ internal class Av1NeighborArrayUnit<T>
 {
     public static readonly T InvalidNeighborData = T.MaxValue;
 
-    [Flags]
-    public enum UnitMask
-    {
-        Left = 1,
-        Top = 2,
-        TopLeft = 4,
-    }
-
     private readonly T[] left;
     private readonly T[] top;
     private readonly T[] topLeft;
@@ -30,6 +22,14 @@ public Av1NeighborArrayUnit(int leftSize, int topSize, int topLeftSize)
         this.topLeft = new T[topLeftSize];
     }
 
+    [Flags]
+    public enum UnitMask
+    {
+        Left = 1,
+        Top = 2,
+        TopLeft = 4,
+    }
+
     public Span<T> Left => this.left;
 
     public Span<T> Top => this.top;

diff --git a/src/ImageSharp/Formats/Heif/Av1/Tiling/Av1PaletteLumaModeInfo.cs b/src/ImageSharp/Formats/Heif/Av1/Tiling/Av1PaletteLumaModeInfo.cs
@@ -1,11 +1,8 @@
-// Copyright (c) Six Labors.
+// Copyright (c) Six Labors.
 // Licensed under the Six Labors Split License.
 
 namespace SixLabors.ImageSharp.Formats.Heif.Av1.Tiling;
 
-internal partial class Av1TileWriter
+internal class Av1PaletteLumaModeInfo
 {
-    internal class Av1PaletteLumaModeInfo
-    {
-    }
 }
diff --git a/src/ImageSharp/Formats/Heif/Av1/Tiling/Av1PictureControlSet.cs b/src/ImageSharp/Formats/Heif/Av1/Tiling/Av1PictureControlSet.cs
@@ -3,20 +3,17 @@
 
 namespace SixLabors.ImageSharp.Formats.Heif.Av1.Tiling;
 
-internal partial class Av1TileWriter
+internal class Av1PictureControlSet
 {
-    internal class Av1PictureControlSet
-    {
-        public required Av1NeighborArrayUnit[] luma_dc_sign_level_coeff_na { get; internal set; }
+    public required Av1NeighborArrayUnit<byte>[] luma_dc_sign_level_coeff_na { get; internal set; }
 
-        public required Av1NeighborArrayUnit[] cr_dc_sign_level_coeff_na { get; internal set; }
+    public required Av1NeighborArrayUnit<byte>[] cr_dc_sign_level_coeff_na { get; internal set; }
 
-        public required Av1NeighborArrayUnit[] cb_dc_sign_level_coeff_na { get; internal set; }
+    public required Av1NeighborArrayUnit<byte>[] cb_dc_sign_level_coeff_na { get; internal set; }
 
-        public required Av1NeighborArrayUnit[] txfm_context_array { get; internal set; }
+    public required Av1NeighborArrayUnit<byte>[] txfm_context_array { get; internal set; }
 
-        public required Av1SequenceControlSet Sequence { get; internal set; }
+    public required Av1SequenceControlSet Sequence { get; internal set; }
 
-        public required Av1PictureParentControlSet Parent { get; internal set; }
-    }
+    public required Av1PictureParentControlSet Parent { get; internal set; }
 }