Quantization: reduce memory footprint without compromising speed and quality

[antonfirsov]

All our current quantizers rely on a very heavy cache in order to produce acceptable quality:

ImageSharp/src/ImageSharp/Processing/Processors/Quantization/EuclideanPixelMap{TPixel}.cs

Line 21 in b06cb32

private readonly ConcurrentDictionary<TPixel, int> distanceCache;

With large images this could consume up to several megabytes of memory every time a larger GIF (or palettized PNG) gets saved. We need to reduce memory usage without compromising speed and quality.

We made a couple of experiments and exchanged many ideas a couple of months ago, but unfortunately all those discussions got lost in the noise of the gitter chatroom. I suggest to recollect and rediscuss those ideas.

Improvements we may consider (according to what I remember from the gitter chat):

1. Improving the RGB Octree implementation to make it very fast by default so we hopefully don't need any cache. Most important: flattening the tree nodes into array(s) of structs (instead of heap objects).
2. Extend the octree to RGBA
3. Replace the dictionary with this cache. It produced promising results in my experiments.
4. New idea from @JimBobSquarePants: replace the dictionary with an LRU cache implementation like the one in BitFaster.Caching.

@saucecontrol since I remember you had really valuable ideas here, I would really appreciate your feedback on all 4 options, and apologies if I'm making you to repeat yourself because of my terrible memory 😄 If I remember correctly, you also had some promising results in your repos.

Bravely assigning milestone 1.1 for now.

[saucecontrol]

apologies if I'm making you to repeat yourself because of my terrible memory

I'm afraid I don't recall much detail of those conversations myself 😄

I did your option 1 for mine, using an octree for the histogram, palette generation, and palette lookup. Unfortunately, due to C#'s lack of first class support for discriminated unions and the JIT's insistence on inserting null checks when dereferencing a member from a struct ptr, the code ended up being quite ugly for perf reasons. But it gives really nice visual results with only about 280KiB of rented pool memory, and in my limited testing was on average an order of magnitude faster than the current ImageSharp octree implementation.

You can find my implementation here: https://github.com/saucecontrol/PhotoSauce/blob/master/src/MagicScaler/Magic/ColorQuantizer.cs

Mine differs quite a bit from any other octree implementation I've seen, which makes it a little difficult to follow, but there's method to any madness you see in there, and I'm happy to answer any questions. Unfortunately I haven't the drawing skills to do good diagrams of how it works, so the code will have to speak for itself for now.

The major architectural differences:

1. I made it GC allocation free by making the octree node a struct and renting a (256KiB) slab of memory from the ArrayPool to hold all the nodes I'll ever have. Once I fill that slab, I reduce the node set by cutting the depth of the tree by one level. Because the tree depth is variable, it has perfect fidelity for low-color images, and it gets faster per-color the more colors it encounters. And when reducing the palette from the possible 8K nodes in the histogram, I weight them based on population count and distance to the parent node color to prioritize the most important colors. This gives better results than simply merging the first available reducible node and allows the palette to have the max possible entries rather than coming up short when all children of a reducible parent are merged.
2. I use a set of LUTs to quickly build the entire octree index for any given input color by interleaving the channel bits, so navigating the tree can be done by simply shifting off 3 bits of that pre-calculated index at a time (I reverse the bit order so the next level index is the low 3 bits). In histogram mode, the tree keeps color values and pixel counts at the node level. In lookup mode, it keeps the palette index and value. I build the octree index in GRB order so that the index roughly approximates luma. When looking for the nearest color within the tree, the nearest populated index at a given level is also going to be the nearest color, provided you search them in the right order. If a node doesn't exist at the minimum leaf level I've defined, I do a full search through the palette and cache the result by creating the missing node.

The LinearBoxesMap idea is clever, but its weakness is of course that it has a pre-defined accuracy limit. If you're trying to fit a 300-color image into a 256-color palette, you should be able to keep perfect matches for the majority of the colors. An octree is great for that because it's a sparse structure, so as long as you make the lookups/navigation fast enough, it does a great job.

The same ideas could be applied to support RGBA in a 'hexadecitree' with 4 bits per node index instead of 3. The alpha channel could be weighted less or quantized in advance. I was only going for GIF support initially, so I haven't yet done any experiments to figure out the best number of nodes, etc.

[antonfirsov]

@saucecontrol thanks for the very detailed reply! What is the pixel format of Span<byte> image ? (Didn't have a chance do dig into your code yet.)

Our big problem is that currently we need the following cache to produce good enough quality when images are transparent:

ImageSharp/src/ImageSharp/Processing/Processors/Quantization/OctreeQuantizer{TPixel}.cs

Lines 114 to 120 in b06cb32

	// Octree only maps the RGB component of a color
	// so cannot tell the difference between a fully transparent
	// pixel and a black one.
	if (this.isDithering || color.Equals(default))
	{
	return (byte)this.pixelMap.GetClosestColor(color, out match);
	}

@JimBobSquarePants in case of pixels having no alpha channel, we likely don't need this path. For users caring about perf we can suggest then to work with Image<Rgb24> all the way.
In the case when alpha channel is present: what if we premultiply with alpha before/while saving the image, instead of using the Euclidean map?

[JimBobSquarePants]

@antonfirsov I found that the difference in quality was noticeable with or without transparency when I experimented with removing the color map entirely.

Looking back at EuclideanPixelMap<TPixel> though I don't know why we use Vector4. Both Octree and Wu convert to Rgba32 when building histograms so we don't need the extra precision. Perhaps we can do something faster there?

[saucecontrol]

What is the pixel format of Span image ?

My implementation works only on BGR24 input. In the final palette derivation step, when merging nodes, I convert the colors to 32bpc linear for better blending/averaging results.

I also have only a single dithering algorithm (a modified Floyd-Steinberg) integrated directly into the palette matching for perf reasons. I didn't really investigate which part of the speed differences between mine and the current ImageSharp implementation come from the histogram vs lookup vs dithering parts of the code, but if your main concern is reducing memory, I do believe you can eliminate the cache without negative perf impact as long as your octree traversal is fast. And of course a faster octree speeds up your histogram step as well so it may be a win-win.

The key to keeping memory usage under control while keeping accuracy is to allow the lookup octree to reach full 8 bit depth for the nodes that contain your palette entries but then restrict the depth for parts of the gamut not covered by the palette. My implementation caps that at 4 bits (a max of 4096 nodes) so I can fit everything in my 256KiB slab. You give up some accuracy on those, but since they represent colors not in the original image anyway (they're used for colors created by dithering error propagation), it's not a big deal.

Again, my use case was animated GIF, so it was very important to me to get it completely optimized for processing hundreds of RGB frames at a go. I don't know if you'd want to generalize your implementation more to allow the same code to handle RGB and RGBA, but if it were me, I'd have separate implementations for those. The RGBA path would naturally require more memory, but my guess would be that around 2MiB would do it if you do some psychovisual trickery with the alpha channel ;)

[antonfirsov]

Looking back at EuclideanPixelMap though I don't know why we use Vector4. Both Octree and Wu convert to Rgba32 when building histograms so we don't need the extra precision. Perhaps we can do something faster there?

ImageSharp/src/ImageSharp/Processing/Processors/Quantization/EuclideanPixelMap{TPixel}.cs

Lines 20 to 21 in b06cb32

	private readonly Vector4[] vectorCache;
	private readonly ConcurrentDictionary<TPixel, int> distanceCache;

@JimBobSquarePants Vector4 is not the problem here since paletteCache has only 256 elements. It's the dictonary that grows out of control. I agree with @saucecontrol that it's the fast Octree itself that should be play the role of the cache and ideally we should eliminate the dictionary, or any similar heavy lookup. Based on #1350 (comment), I think MagicScaler's Octree should provide better results in terms of quality than ours. For me this sounds like something worth a POC as a next step. @JimBobSquarePants you still skeptical? Do you see any issues I don't? (Especially re dithering.)

Note: converting pixel spans to Bgr24 should be cheap after implementing #1354

[JimBobSquarePants]

@antonfirsov I only mention the Vector4 usage because converting to/from RGB compatible types is much faster.

My biggest concern here is that we are focusing too much on Octree. All of the quantizers use the pixel map to improve output quality (due to naïve color distance formula) and two of the quantizers depend on it entirely for palette application.

ImageSharp/src/ImageSharp/Processing/Processors/Quantization/PaletteQuantizer{TPixel}.cs

Lines 63 to 64 in b06cb32

	public readonly byte GetQuantizedColor(TPixel color, out TPixel match)
	=> (byte)this.pixelMap.GetClosestColor(color, out match);

So... Perhaps we can do something to allow fast linear color space conversion using LUTs since we're working with pixel formats with byte accuracy in the quantizers @saucecontrol I know you do this but I can never read your code (I'm convinced you've been sent here from the future) which could improve the accuracy of the built in palette matching mechanisms. Maybe if the quality is improved we can drop the pixel map entirely.

Then we can figure out what to do for the other quantizers that depend on it.

The pixel map was always a crutch. That's why I buried it far away from the API surface.

[saucecontrol]

@JimBobSquarePants I think I might be confused about how your implementation works now. I thought the purpose of the pixel map was to improve speed by eliminating the need for a full search through the palette to find the nearest palette color for each input pixel color. You're saying it has something to do with improving quality? How does that work?

You can actually use an octree for palette lookup regardless of how the palette was originally obtained, so the idea isn't necessarily that you go all octree all the time. If you needed to map an image to e.g. a standard web palette, you would build out an octree with the values from the palette to start. As you map the input image to the palette, any time you navigate the octree and the node for your pixel color doesn't exist, you do the linear search through the palette and then cache the winner by creating that node so that next time you see that color or one like it, it's already there. Essentially, the octree works as a sparse cache in that mode. The real win is in the fact that it can group similar colors instead of having a key per distinct input value, because you pick the max tree depth for those cache nodes. It's actually very similar to @antonfirsov's LinearBoxesMap concept except that the accuracy is variable (from 4 to 8 bits in my implementation) instead of being fixed at e.g. 5 bits, and any unused part of the color space never gets allocated any memory space.

[JimBobSquarePants]

@saucecontrol Using the pixel map actually results in much slower quantization than not using one.

The quality improvement is achieved by finding the closest color in the resolved palette to the given pixel via Euclidian distance. I originally started using that approach because a bug in my dithering caused exceptions when trying the match the dithered result against the palette. I stick with it now that issue is gone because I still see quite dramatic differences in quality when compared against both Wu and Octree quantizers.

The dictionary caches results since the distance calculation is horribly slow.

ImageSharp/src/ImageSharp/Processing/Processors/Quantization/EuclideanPixelMap{TPixel}.cs

Lines 72 to 103 in b06cb32

	private int GetClosestColorSlow(TPixel color, ref TPixel paletteRef, out TPixel match)
	{
	// Loop through the palette and find the nearest match.
	int index = 0;
	float leastDistance = float.MaxValue;
	var vector = color.ToVector4();
	ref Vector4 vectorCacheRef = ref MemoryMarshal.GetReference<Vector4>(this.vectorCache);
	for (int i = 0; i < this.Palette.Length; i++)
	{
	Vector4 candidate = Unsafe.Add(ref vectorCacheRef, i);
	float distance = Vector4.DistanceSquared(vector, candidate);
	// If it's an exact match, exit the loop
	if (distance == 0)
	{
	index = i;
	break;
	}
	if (distance < leastDistance)
	{
	// Less than... assign.
	index = i;
	leastDistance = distance;
	}
	}
	// Now I have the index, pop it into the cache for next time
	this.distanceCache[color] = index;
	match = Unsafe.Add(ref paletteRef, index);
	return index;
	}

I'm intrigued by the variable bit depth approach. From our previous conversations the output of your quantizer color wise was amazing (though you were having dithering issues)

[antonfirsov]

All of the quantizers use the pixel map to improve output quality

This is the problem. We need a significantly more memory-efficient lookup mechanism.

Octree comes in the picture as probably the best candidate for such a lookup in case of 3-component images, assuming that @saucecontrol 's implementation is producing decent quality, and it's possible to integrate it with our generic dithering logic.

[saucecontrol]

Ah, yeah, we're talking about the same thing. I think @antonfirsov and I were phrasing it poorly by talking about eliminating the cache. The map is a necessity as long as you're doing dithering, because even if you had a perfect histogram to start (which you can't always), error diffusion might still introduce colors that fall outside the input image's gamut. And you'll want that map to cache any new palette entries it has resolved.

The real issue is that a dictionary isn't a good map/cache because its keys are opaque integers and its allocations are GC-heavy. Using an octree as the map/cache can potentially solve both of those.

I don't know if there's any part of my code you can lift directly, just because it's so tuned for speed specific to my use case. Reviewing it now, there's some ugly stuff in there. Conceptually, though, I think you could model after it because I haven't seen anything that balances quality, speed, and memory as well it does.

I did end up resolving my dithering issues we chatted about on gitter, by the way. My final implementation just ended up being Floyd-Steinberg with the error weighted to 7/8. I'd say the quality is generally equal or better than what ImageSharp does today, although there's not a ton of improvement to be had there -- your implementation is already very good quality-wise.

[JimBobSquarePants]

The real issue is that a dictionary isn't a good map/cache because its keys are opaque integers and its allocations are GC-heavy. Using an octree as the map/cache can potentially solve both of those.

Yep, that's why I was going to just insert an LRU cache in there for now but if there's a better cache then I'm happy for anyone to have a go at replacing it.

If there was a cheaper way to match the dithered result to the palette and retain (or improve) output quality then I'm all ears. It'll need someone else to experiment/implement though.

[antonfirsov]

I don't think a dictonary-based LRU cache is a good candidate, it's still very heavy, and at the same time likely much slower than a dictionary. In other words: it seems like a low hanging fruit, but I can bet for 95% that it just won't work well enough, so I would discourage investing into it.

If we want a "best effort" map as a quick workaround, let's use LinearBoxesMap. It can cache up to 4 colors per 5-byte box for exact match, just like a dictionary. I made some experiments to investigate the likehood of colors overflowing boxes, and for most palettes it was very unlikely. In those rare cases when a color does not fit, the map will return "misses" and the linear euclidean search path activates (just like it would happen with LRU). For images with less uniform palette distribution (eg. monochrome or monocolor images), the result will be much worse of course.