Complete core modular functions

src/ImageSharp/Formats/Jxl/Processing/Modular/Encoding/JxlModularEncoding.cs

@@ -33,7 +33,7 @@ internal static class JxlModularEncoding
                 return false;
             }
 
-            JxlMaNode node = tree[cur.Pos];
+            JxlFlatDecisionNode node = tree[cur.Pos];
 
             // Leaf.
             if (node.Property0 == -1)
@@ -76,7 +76,7 @@ internal static class JxlModularEncoding
                 continue;
             }
 
-            if (node.Properties[0] >= NumStaticProperties)
+            if (node.Properties[0] >= JxlPredictorFacts.StaticProperties)
             {
                 ranges.Push(new(node.SplitValues[0], cur.End, node.ChildID));
                 ranges.Push(new(node.SplitValue0, node.SplitValues[0], node.ChildID + 1));
@@ -87,7 +87,7 @@ internal static class JxlModularEncoding
             }
 
             // <= side
-            if (node.Properties[1] >= NumStaticProperties)
+            if (node.Properties[1] >= JxlPredictorFacts.StaticProperties)
             {
                 ranges.Push(
                     new(node.SplitValues[1], node.SplitValue0, node.ChildID + 2));
@@ -116,12 +116,12 @@ internal static class JxlModularEncoding
             {
                 hasWeightedPrediction = true;
             }
-            else if (p >= NumStaticProperties)
+            else if (p >= JxlPredictorFacts.StaticProperties)
             {
                 hasNonWeighted = true;
             }
 
-            if (p >= NumStaticProperties && p != GradientProp)
+            if (p >= JxlPredictorFacts.StaticProperties && p != GradientProp)
             {
                 gradientOnly = false;
             }
@@ -136,7 +136,7 @@ internal static class JxlModularEncoding
             int cur = nodes.Peek();
             _ = nodes.Dequeue();
 
-            while (globalTree[cur].Property < NumStaticProperties && globalTree[cur].Property != -1)
+            while (globalTree[cur].Property is < JxlPredictorFacts.StaticProperties and not -1)
             {
                 if (staticProps[globalTree[cur].Property] > globalTree[cur].SplitValue)
                 {
@@ -176,7 +176,7 @@ internal static class JxlModularEncoding
             {
                 int currentChild = i == 0 ? globalTree[cur].LeftChild : globalTree[cur].RightChild;
 
-                while (globalTree[currentChild].Property < kNumStaticProperties && globalTree[currentChild].Property != -1)
+                while (globalTree[currentChild].Property < JxlPredictorFacts.StaticProperties && globalTree[currentChild].Property != -1)
                 {
                     if (staticProps[globalTree[currentChild].Property] >
                         globalTree[currentChild].SplitValue)
@@ -220,7 +220,7 @@ internal static class JxlModularEncoding
 
         if (numProps > JxlMaConstants.NumTreeContexts)
         {
-            numProps = (JxlMath.DivCeil(numProps - JxlMaConstants.NumTreeContexts, kExtraPropsPerChannel) * kExtraPropsPerChannel) + JxlMaConstants.NumTreeContexts;
+            numProps = (JxlMath.DivCeil(numProps - JxlMaConstants.NumTreeContexts, JxlContextPrediction.ExtraPropertiesPerChannel) * JxlContextPrediction.ExtraPropertiesPerChannel) + JxlMaConstants.NumTreeContexts;
         }
         else
         {
@@ -281,7 +281,7 @@ internal static class JxlModularEncoding
                         return false;
                     }
                 }
-                else if (n.Property >= NumStaticProperties)
+                else if (n.Property >= JxlPredictorFacts.StaticProperties)
                 {
                     return false;
                 }

src/ImageSharp/Formats/Jxl/Processing/Modular/Encoding/JxlModularEncodingEncoder.cs

@@ -264,11 +264,11 @@ internal static class JxlModularEncodingEncoder
 
         treeSamples.AllSamplesDone();
 
-        JxlMaEncoder.ComputeBestTree(treeSamples, options.SplittingHeuristicsModeThreshold * requiredCost, info, staticPropertyRange, options.FastDecodeMultiplier, tree);
+        JxlMaEncoder.ComputeBestTree(treeSamples, options.SplittingHeuristicsNodeThreshold * requiredCost, info, staticPropertyRange, options.FastDecodeMultiplier, tree);
         return tree;
     }
 
-    public static void EncodeModularChannelMAANS(JxlModularImage image, int channelIndex, JxlModularHeader wpHeader, Tree globalTree, Span<JxlToken> tokens, int groupId, bool skipEncoderFastPath)
+    public static void EncodeModularChannelMAANS(Configuration configuration, JxlModularImage image, int channelIndex, JxlModularHeader wpHeader, Tree globalTree, Span<JxlToken> tokens, int groupId, bool skipEncoderFastPath)
     {
         JxlModularChannel channel = image.Channels[channelIndex];
 
@@ -563,7 +563,7 @@ internal static class JxlModularEncodingEncoder
         throw new InvalidOperationException("Invalid tree type: " + treeKind);
     }
 
-    public static Tree LearnTree(JxlModularImage[] images, Span<JxlModularOptions> options, int start, int stop, List<JxlModularMultiplierInfo>? multiplierInfo = null)
+    public static Tree LearnTree(Configuration configuration, JxlModularImage[] images, Span<JxlModularOptions> options, int start, int stop, List<JxlModularMultiplierInfo>? multiplierInfo = null)
     {
         multiplierInfo ??= [];
 
@@ -581,7 +581,7 @@ internal static class JxlModularEncodingEncoder
         for (int i = start; i < stop; i++)
         {
             maxC = Math.Max(images[i].Channels.Count, maxC);
-            JxlMaEncoder.CollectPixelSamples(images[i], options[i], i, groupPixelCount, channelPixelCount, pixelSamples, diffSamples);
+            JxlMaEncoder.CollectPixelSamples(configuration, images[i], options[i], i, groupPixelCount, channelPixelCount, pixelSamples, diffSamples);
         }
 
         // StaticPropRange range;
@@ -596,7 +596,7 @@ internal static class JxlModularEncodingEncoder
         currRange[1] = stop;
         range[1] = currRange;
 
-        samples.PreQuantizeProperties(range, multiplierInfo, groupPixelCount, channelPixelCount, pixelSamples, diffSamples, options[start].MaxPropertyValues);
+        samples.PreQuantizeProperties(configuration, range, multiplierInfo, groupPixelCount, channelPixelCount, pixelSamples, diffSamples, options[start].MaxPropertyValues);
 
         int totalPixels = 0;
 
@@ -618,9 +618,9 @@ internal static class JxlModularEncodingEncoder
 
         for (int i = start; i < stop; i++)
         {
-            int nb_channels = images[i].Channels.Count;
+            int numChannels = images[i].Channels.Count;
 
-            if (images[i].Width == 0 || images[i].Height == 0 || nb_channels < 1)
+            if (images[i].Width == 0 || images[i].Height == 0 || numChannels < 1)
             {
                 continue;
             }
@@ -642,7 +642,7 @@ internal static class JxlModularEncodingEncoder
                 JxlContextPrediction.SetPredictorMode(options[i].WpMode, wpHeader);
             }
 
-            for (int c = 0; c < nb_channels; c++)
+            for (int c = 0; c < numChannels; c++)
             {
                 if (c >= images[i].MetaChannels &&
                     (images[i].Channels[c].Width > options[i].MaxChannelSize ||
@@ -656,7 +656,7 @@ internal static class JxlModularEncodingEncoder
                     continue;  // skip empty channels
                 }
 
-                GatherTreeData(images[i], c, i, wpHeader, options[i], samples, totalPixels);
+                GatherTreeData(configuration, images[i], c, i, wpHeader, options[i], samples, totalPixels);
             }
         }
 
@@ -724,7 +724,7 @@ internal static class JxlModularEncodingEncoder
             {
                 if (i >= image.MetaChannels &&
                     (image.Channels[i].Width > options.MaxChannelSize ||
-                     image.Channels[i].Height > options.v))
+                     image.Channels[i].Height > options.MaxChannelSize))
                 {
                     break;
                 }
@@ -734,7 +734,7 @@ internal static class JxlModularEncodingEncoder
                     continue;
                 }
 
-                EncodeModularChannelMAANS(image, i, header.WeightedHeader, tree, CollectionsMarshal.AsSpan(tokens)[tokenp..], groupId, options.SkipEncoderFastPath);
+                EncodeModularChannelMAANS(configuration, image, i, header.WeightedHeader, tree, CollectionsMarshal.AsSpan(tokens)[tokenp..], groupId, options.SkipEncoderFastPath);
             }
 
             if (tokenp != tokens.Count)
@@ -746,7 +746,7 @@ internal static class JxlModularEncodingEncoder
         width = imageWidth;
     }
 
-    public static void ModularGenericCompress(JxlModularImage image, JxlModularOptions options, JxlBitWriter writer, JxlAuxiliaryOutput auxOut, JxlLayerType layerType, int groupId)
+    public static void ModularGenericCompress(Configuration configuration, JxlModularImage image, JxlModularOptions options, JxlBitWriter writer, JxlAuxiliaryOutput auxOut, JxlLayerType layerType, int groupId)
     {
         int numChannels = image.Channels.Count;
 
@@ -785,7 +785,7 @@ internal static class JxlModularEncodingEncoder
 
         if (modularOptions.TreeKind == JxlTreeKind.Learn)
         {
-            tree = LearnTree(image, options, 0, 1);
+            tree = LearnTree(configuration, image, options, 0, 1);
         }
         else
         {

src/ImageSharp/Formats/Jxl/Processing/Modular/JxlModularChannel.cs

@@ -8,7 +8,7 @@ namespace SixLabors.ImageSharp.Formats.Jxl.Processing.Modular;
 /// <summary>
 /// A wrapper over <see cref="JxlImageI"/> for modular operations.
 /// </summary>
-internal sealed class JxlModularChannel
+internal sealed class JxlModularChannel : IDisposable
 {
     public JxlModularChannel(Configuration configuration, int width, int height, int horizShift, int vertShift)
     {
@@ -78,4 +78,6 @@ internal sealed class JxlModularChannel
     public Span<int> GetRowMinus(int y, int minus) => this.Plane.GetRowMinus(y, minus);
 
     public Span<int> GetRowPlus(int y, int plus) => this.Plane.GetRowPlus(y, plus);
+
+    public void Dispose() => this.Plane.Dispose();
 }

src/ImageSharp/Formats/Jxl/Processing/Modular/JxlModularImage.cs

@@ -1,12 +1,31 @@
 // Copyright (c) Six Labors.
 // Licensed under the Six Labors Split License.
 
+using SixLabors.ImageSharp.Formats.Jxl.Processing.Modular.Encoding.ContextPrediction;
+using SixLabors.ImageSharp.Formats.Jxl.Processing.Modular.Transforms;
+
 namespace SixLabors.ImageSharp.Formats.Jxl.Processing.Modular;
 
-internal sealed class JxlModularImage
+internal sealed class JxlModularImage : IDisposable
 {
+    public JxlModularImage(int width, int height, int bitDepth)
+    {
+        this.Width = width;
+        this.Height = height;
+        this.BitDepth = bitDepth;
+        this.MetaChannels = 0;
+        this.IsInvalid = false;
+    }
+
+    public JxlModularImage()
+        : this(0, 0, 8)
+    {
+    }
+
     public List<JxlModularChannel> Channels { get; set; } = [];
 
+    public List<JxlTransform> Transforms { get; set; } = [];
+
     /// <summary>
     /// Gets or sets the total number of metachannels in this image.
     /// </summary>
@@ -16,4 +35,56 @@ internal sealed class JxlModularImage
     /// Gets or sets the bit depth used in this image.
     /// </summary>
     public int BitDepth { get; set; }
+
+    /// <summary>
+    /// Gets or sets modular image width.
+    /// </summary>
+    public int Width { get; set; }
+
+    /// <summary>
+    /// Gets or sets modular image height.
+    /// </summary>
+    public int Height { get; set; }
+
+    /// <summary>
+    /// Gets or sets a value indicating whether the modular image has an error.
+    /// </summary>
+    public bool IsInvalid { get; set; }
+
+    public bool IsEmpty => this.Channels.Any(x => x.Width > 0 && x.Height > 0);
+
+    public void UndoTransforms(Configuration configuration, JxlModularHeader wpHeader)
+    {
+        while (this.Transforms.Count > 0)
+        {
+            JxlTransform transform = this.Transforms.First();
+            transform.Inverse(configuration, this, wpHeader);
+            this.Transforms.RemoveAt(0);
+        }
+    }
+
+    public static JxlModularImage Create(Configuration configuration, int width, int height, int bitDepth, int channels)
+    {
+        JxlModularImage result = new(width, height, bitDepth);
+
+        for (int i = 0; i < channels; i++)
+        {
+            JxlModularChannel c = new(configuration, width, height, 0, 0);
+
+            result.Channels.Add(c);
+            result.Channels[^1].Component = i;
+        }
+
+        return result;
+    }
+
+    public void Dispose()
+    {
+        foreach (JxlModularChannel channel in this.Channels)
+        {
+            channel.Dispose();
+        }
+
+        this.Channels = [];
+    }
 }

src/ImageSharp/Formats/Jxl/Processing/Modular/JxlModularOptions.cs

@@ -0,0 +1,64 @@
+// Copyright (c) Six Labors.
+// Licensed under the Six Labors Split License.
+
+using SixLabors.ImageSharp.Formats.Jxl.Processing.Encoder.Ans;
+using SixLabors.ImageSharp.Formats.Jxl.Processing.Modular.Encoding.ContextPrediction;
+
+namespace SixLabors.ImageSharp.Formats.Jxl.Processing.Modular;
+
+internal sealed class JxlModularOptions
+{
+    /// <summary>
+    /// Gets or sets value indicating when to stop decoding/encoding
+    /// when reaching a non-meta channel that has a dimension bigger
+    /// than this value.
+    /// </summary>
+    public int MaxChannelSize { get; set; } = 0xFFFFFF;
+
+    /// <summary>
+    /// Gets or sets value used during decoding for validation
+    /// of transforms (squeezing) scheme.
+    /// </summary>
+    public int GroupDimension { get; set; } = 0x1FFFFFFF;
+
+    /// <summary>
+    /// Gets or sets fraction of pixels to look at to learn a MA tree.
+    /// </summary>
+    public float NumberOfRepeats { get; set; } = 0.5f;
+
+    /// <summary>
+    /// Gets or sets maximum number of previous channel properties
+    /// to use in the MA trees.
+    /// </summary>
+    public int MaxProperties { get; set; }
+
+    /// <summary>
+    /// Gets or sets properties that default to channel, group, weighted,
+    /// gradient residual, W-NW, NW-N, N-NE, N-NN.
+    /// </summary>
+    public List<int> SplittingHeuristicsProperties { get; set; } = [0, 1, 15, 9, 10, 11, 12, 13];
+
+    public float SplittingHeuristicsNodeThreshold { get; set; } = 96f;
+
+    public int MaxPropertyValues { get; set; } = 32;
+
+    public JxlPredictor Predictor { get; set; } = JxlPredictor.Undefined;
+
+    public int WpMode { get; set; }
+
+    public float FastDecodeMultiplier { get; set; } = 1.01f;
+
+    public JxlTreeMode WpTreeMode { get; set; } = JxlTreeMode.Default;
+
+    public bool SkipEncoderFastPath { get; set; }
+
+    public JxlTreeKind TreeKind { get; set; } = JxlTreeKind.Learn;
+
+    public JxlHistogramParameters HistogramParameters { get; set; } = new();
+
+    /// <summary>
+    /// Gets or sets a value indicating whether to ignore the image and just
+    /// pretend all tokens are zeroes.
+    /// </summary>
+    public bool ZeroTokens { get; set; }
+}

src/ImageSharp/Formats/Jxl/Processing/Modular/Transforms/JxlTransform.cs

@@ -2,12 +2,47 @@
 // Licensed under the Six Labors Split License.
 
 using System.Numerics.Tensors;
+using System.Runtime.InteropServices;
 using SixLabors.ImageSharp.Formats.Jxl.Fields;
+using SixLabors.ImageSharp.Formats.Jxl.Processing.Modular.Encoding.ContextPrediction;
 
 namespace SixLabors.ImageSharp.Formats.Jxl.Processing.Modular.Transforms;
 
 internal sealed class JxlTransform : IJxlFields
 {
+    public JxlTransform()
+        : this(JxlTransformType.Invalid)
+    {
+    }
+
+    public JxlTransform(JxlTransformType id)
+    {
+        JxlBundle.Init(this);
+        this.TransformType = id;
+    }
+
+    public JxlTransformType TransformType { get; set; }
+
+    public int BeginC { get; set; }
+
+    public int RctType { get; set; }
+
+    public int NumC { get; set; }
+
+    public int Colors { get; set; }
+
+    public int Deltas { get; set; }
+
+    public List<JxlSqueezeParameters> Squeezes { get; set; } = [];
+
+    public int MaxDeltaError { get; set; }
+
+    public JxlPredictor Predictor { get; set; }
+
+    public bool OrderedPalette { get; set; } = true;
+
+    public bool LossyPalette { get; set; }
+
     public bool Visit(JxlVisitor visitor) => throw new NotImplementedException();
 
     public static void CheckEqualChannels(JxlModularImage image, int c1, int c2)
@@ -19,7 +54,7 @@ internal sealed class JxlTransform : IJxlFields
             throw new InvalidOperationException($"Invalid channel range: {c1}..{c2} (there are only {channelsCount} channels)");
         }
 
-        if (c1 < image.NbMetaChannels && c2 >= image.NbMetaChannels)
+        if (c1 < image.MetaChannels && c2 >= image.MetaChannels)
         {
             throw new InvalidOperationException("Invalid: transforming mix of meta and nonmeta");
         }
@@ -56,4 +91,46 @@ internal sealed class JxlTransform : IJxlFields
             max = Math.Max(maxRow, max);
         }
     }
+
+    public void Inverse(Configuration configuration, JxlModularImage input, JxlModularHeader wpHeader)
+    {
+        switch (this.TransformType)
+        {
+            case JxlTransformType.Rct:
+                JxlRct.InverseRct(configuration, input, this.BeginC, this.RctType);
+                break;
+
+            case JxlTransformType.Squeeze:
+                JxlSqueeze.InverseSqueeze(configuration, input, CollectionsMarshal.AsSpan(this.Squeezes));
+                break;
+
+            case JxlTransformType.Palette:
+                JxlPalette.InversePalette(configuration, input, this.BeginC, this.Colors, this.Deltas, this.Predictor, wpHeader);
+                break;
+
+            default:
+                throw new InvalidOperationException($"Unknown transform: {this.TransformType}");
+        }
+    }
+
+    public void MetaApply(Configuration configuration, JxlModularImage image)
+    {
+        switch (this.TransformType)
+        {
+            case JxlTransformType.Rct:
+                CheckEqualChannels(image, this.BeginC, this.BeginC + 2);
+                break;
+
+            case JxlTransformType.Squeeze:
+                JxlSqueeze.MetaSqueeze(configuration, image, this.Squeezes);
+                break;
+
+            case JxlTransformType.Palette:
+                JxlPalette.MetaPalette(configuration, image, this.BeginC, this.BeginC + this.NumC - 1, this.Colors, this.Deltas);
+                break;
+
+            default:
+                throw new InvalidOperationException($"Unknown transform: {this.TransformType}");
+        }
+    }
 }