From 1ba8a6f80f794d6973b505ba63f415dd73b6e6b7 Mon Sep 17 00:00:00 2001 From: Lucas CHOLLET Date: Tue, 14 May 2024 05:24:33 -0400 Subject: [PATCH] LibGfx: Add an implementation of the MedianCut algorithm This is useful to find the best matching color palette from an existing bitmap. It can be used in PixelPaint but also in encoders of old image formats that only support indexed colors e.g. GIF. --- Tests/LibGfx/CMakeLists.txt | 1 + Tests/LibGfx/TestMedianCut.cpp | 56 +++++++ Userland/Libraries/LibGfx/CMakeLists.txt | 1 + Userland/Libraries/LibGfx/MedianCut.cpp | 185 +++++++++++++++++++++++ Userland/Libraries/LibGfx/MedianCut.h | 53 +++++++ 5 files changed, 296 insertions(+) create mode 100644 Tests/LibGfx/TestMedianCut.cpp create mode 100644 Userland/Libraries/LibGfx/MedianCut.cpp create mode 100644 Userland/Libraries/LibGfx/MedianCut.h diff --git a/Tests/LibGfx/CMakeLists.txt b/Tests/LibGfx/CMakeLists.txt index b0705db1665..2c5e42c6a57 100644 --- a/Tests/LibGfx/CMakeLists.txt +++ b/Tests/LibGfx/CMakeLists.txt @@ -8,6 +8,7 @@ set(TEST_SOURCES TestICCProfile.cpp TestImageDecoder.cpp TestImageWriter.cpp + TestMedianCut.cpp TestPainter.cpp TestParseISOBMFF.cpp TestRect.cpp diff --git a/Tests/LibGfx/TestMedianCut.cpp b/Tests/LibGfx/TestMedianCut.cpp new file mode 100644 index 00000000000..f116ae82438 --- /dev/null +++ b/Tests/LibGfx/TestMedianCut.cpp @@ -0,0 +1,56 @@ +/* + * Copyright (c) 2024, Lucas Chollet + * + * SPDX-License-Identifier: BSD-2-Clause + */ + +#include +#include +#include + +TEST_CASE(single_element) +{ + auto const bitmap = TRY_OR_FAIL(Gfx::Bitmap::create(Gfx::BitmapFormat::BGRA8888, { 1, 1 })); + bitmap->set_pixel(0, 0, Gfx::Color::NamedColor::White); + + auto const result = TRY_OR_FAIL(Gfx::median_cut(bitmap, 1)); + + EXPECT_EQ(result.palette().size(), 1ul); + EXPECT_EQ(result.closest_color(Gfx::Color::NamedColor::White), Gfx::Color::NamedColor::White); +} + +namespace { +constexpr auto colors = to_array({ { 253, 0, 0 }, { 255, 0, 0 }, { 0, 253, 0 }, { 0, 255, 0 } }); + +ErrorOr> create_test_bitmap() +{ + auto bitmap = TRY(Gfx::Bitmap::create(Gfx::BitmapFormat::BGRA8888, { colors.size(), 1 })); + for (u8 i = 0; i < colors.size(); ++i) + bitmap->set_pixel(i, 0, colors[i]); + return bitmap; +} +} + +TEST_CASE(four_in_four_out) +{ + auto const bitmap = TRY_OR_FAIL(create_test_bitmap()); + + auto const result = TRY_OR_FAIL(Gfx::median_cut(bitmap, 4)); + + EXPECT_EQ(result.palette().size(), 4ul); + for (auto const color : colors) + EXPECT_EQ(result.closest_color(color), color); +} + +TEST_CASE(four_in_two_out) +{ + auto const bitmap = TRY_OR_FAIL(create_test_bitmap()); + + auto const result = TRY_OR_FAIL(Gfx::median_cut(bitmap, 2)); + + EXPECT_EQ(result.palette().size(), 2ul); + EXPECT_EQ(result.closest_color(Gfx::Color(253, 0, 0)), Gfx::Color(254, 0, 0)); + EXPECT_EQ(result.closest_color(Gfx::Color(255, 0, 0)), Gfx::Color(254, 0, 0)); + EXPECT_EQ(result.closest_color(Gfx::Color(0, 253, 0)), Gfx::Color(0, 254, 0)); + EXPECT_EQ(result.closest_color(Gfx::Color(0, 255, 0)), Gfx::Color(0, 254, 0)); +} diff --git a/Userland/Libraries/LibGfx/CMakeLists.txt b/Userland/Libraries/LibGfx/CMakeLists.txt index 96c26405ba2..3ae75f02423 100644 --- a/Userland/Libraries/LibGfx/CMakeLists.txt +++ b/Userland/Libraries/LibGfx/CMakeLists.txt @@ -74,6 +74,7 @@ set(SOURCES ImageFormats/WebPWriter.cpp ImageFormats/WebPWriterLossless.cpp ImmutableBitmap.cpp + MedianCut.cpp Painter.cpp Palette.cpp Path.cpp diff --git a/Userland/Libraries/LibGfx/MedianCut.cpp b/Userland/Libraries/LibGfx/MedianCut.cpp new file mode 100644 index 00000000000..1f7b362e1c2 --- /dev/null +++ b/Userland/Libraries/LibGfx/MedianCut.cpp @@ -0,0 +1,185 @@ +/* + * Copyright (c) 2024, Lucas Chollet + * + * SPDX-License-Identifier: BSD-2-Clause + */ + +#include +#include +#include + +namespace Gfx { + +namespace { + +using Bucket = Vector; +using Buckets = Vector; + +void sort_along_color(Bucket& bucket, u8 color_index) +{ + auto less_than = [=](ARGB32 first, ARGB32 second) { + auto const first_color = Color::from_argb(first); + auto const second_color = Color::from_argb(second); + switch (color_index) { + case 0: + return first_color.red() < second_color.red(); + case 1: + return first_color.green() < second_color.green(); + case 2: + return first_color.blue() < second_color.blue(); + default: + VERIFY_NOT_REACHED(); + } + }; + + AK::quick_sort(bucket, less_than); +} + +template +struct MaxAndIndex { + T maximum; + u32 index; +}; + +template +MaxAndIndex max_and_index(Span values, GreaterThan greater_than) +{ + VERIFY(values.size() != 0); + + u32 max_index = 0; + RemoveCV max_value = values[0]; + for (u32 i = 0; i < values.size(); ++i) { + if (greater_than(values[i], max_value)) { + max_value = values[i]; + max_index = i; + } + } + + return { max_value, max_index }; +} + +ErrorOr split_bucket(Buckets& buckets, u32 index_to_split_at, u8 color_index) +{ + auto& to_split = buckets[index_to_split_at]; + + sort_along_color(to_split, color_index); + + Bucket new_bucket {}; + + auto const middle = to_split.size() / 2; + + auto const span_to_move = to_split.span().slice(middle); + // FIXME: Make Vector::try_extend() take a span + for (u32 i = 0; i < span_to_move.size(); ++i) + TRY(new_bucket.try_append(span_to_move[i])); + to_split.remove(middle, span_to_move.size()); + + TRY(buckets.try_append(move(new_bucket))); + + return {}; +} + +struct IndexAndChannel { + u32 bucket_index {}; + float score {}; + u8 color_index {}; +}; + +ErrorOr> find_largest_bucket(Buckets const& buckets) +{ + Vector bucket_stats {}; + + for (u32 i = 0; i < buckets.size(); ++i) { + auto const& bucket = buckets[i]; + + if (bucket.size() == 1) + continue; + + Statistics red {}; + Statistics green {}; + Statistics blue {}; + for (auto const argb : bucket) { + auto const color = Color::from_argb(argb); + red.add(color.red()); + green.add(color.green()); + blue.add(color.blue()); + } + + Array const variances = { red.variance(), green.variance(), blue.variance() }; + + auto const stats = max_and_index(variances.span(), [](auto a, auto b) { return a > b; }); + + TRY(bucket_stats.try_append({ i, stats.maximum, static_cast(stats.index) })); + } + + if (bucket_stats.size() == 0) + return OptionalNone {}; + + return bucket_stats[max_and_index(bucket_stats.span(), [](auto a, auto b) { return a.score > b.score; }).index]; +} + +ErrorOr split_largest_bucket(Buckets& buckets) +{ + if (auto const bucket_info = TRY(find_largest_bucket(buckets)); bucket_info.has_value()) + TRY(split_bucket(buckets, bucket_info->bucket_index, bucket_info->color_index)); + + return {}; +} + +ErrorOr color_palette_from_buckets(Buckets const& buckets) +{ + Vector palette; + HashMap conversion_table; + + for (auto const& bucket : buckets) { + u32 average_r {}; + u32 average_g {}; + u32 average_b {}; + + for (auto const argb : bucket) { + auto const color = Color::from_argb(argb); + average_r += color.red(); + average_g += color.green(); + average_b += color.blue(); + } + + auto const bucket_size = bucket.size(); + auto const average_color = Color( + round_to(static_cast(average_r) / bucket_size), + round_to(static_cast(average_g) / bucket_size), + round_to(static_cast(average_b) / bucket_size)); + + TRY(palette.try_append(average_color)); + for (auto const color : bucket) + TRY(conversion_table.try_set(Color::from_argb(color), { average_color, palette.size() - 1 })); + } + + return ColorPalette { move(palette), move(conversion_table) }; +} + +} + +ErrorOr median_cut(Bitmap const& bitmap, u16 palette_size) +{ + HashTable color_set; + for (auto color : bitmap) + TRY(color_set.try_set(color)); + + Vector first_bucket; + TRY(first_bucket.try_ensure_capacity(color_set.size())); + for (auto const color : color_set) + first_bucket.append(color); + + Buckets bucket_list; + TRY(bucket_list.try_append(first_bucket)); + + u16 old_bucket_size = 0; + while (bucket_list.size() > old_bucket_size && bucket_list.size() < palette_size) { + old_bucket_size = bucket_list.size(); + TRY(split_largest_bucket(bucket_list)); + } + + return color_palette_from_buckets(bucket_list); +} + +} diff --git a/Userland/Libraries/LibGfx/MedianCut.h b/Userland/Libraries/LibGfx/MedianCut.h new file mode 100644 index 00000000000..771004848ae --- /dev/null +++ b/Userland/Libraries/LibGfx/MedianCut.h @@ -0,0 +1,53 @@ +/* + * Copyright (c) 2024, Lucas Chollet + * + * SPDX-License-Identifier: BSD-2-Clause + */ + +#pragma once + +#include +#include +#include +#include + +namespace Gfx { + +class ColorPalette { +public: + struct ColorAndIndex { + Color color; + size_t index; + }; + + ColorPalette(Vector palette, HashMap conversion_table) + : m_palette(move(palette)) + , m_conversion_table(move(conversion_table)) + { + } + + Vector const& palette() const + { + return m_palette; + } + + Color closest_color(Color input) const + { + return m_palette[index_of_closest_color(input)]; + } + + u32 index_of_closest_color(Color input) const + { + if (auto const result = m_conversion_table.get(input); result.has_value()) + return result->index; + TODO(); + } + +private: + Vector m_palette; + HashMap m_conversion_table; +}; + +ErrorOr median_cut(Bitmap const& bitmap, u16 palette_size); + +}