diff --git a/Userland/Libraries/LibGfx/ImageFormats/JPEGWriter.cpp b/Userland/Libraries/LibGfx/ImageFormats/JPEGWriter.cpp index d97171343109ec..a9ea97f429588f 100644 --- a/Userland/Libraries/LibGfx/ImageFormats/JPEGWriter.cpp +++ b/Userland/Libraries/LibGfx/ImageFormats/JPEGWriter.cpp @@ -9,6 +9,7 @@ #include "JPEGWriterTables.h" #include #include +#include #include #include #include @@ -256,39 +257,35 @@ class JPEGEncodingContext { } } - ErrorOr write_huffman_stream(Mode mode) + ErrorOr huffman_encode_macroblocks(Mode mode) { for (auto& float_macroblock : m_macroblocks) { auto macroblock = float_macroblock.as_i16(); - TRY(encode_dc(dc_luminance_huffman_table, macroblock.y, 0)); - TRY(encode_ac(ac_luminance_huffman_table, macroblock.y)); - - TRY(encode_dc(dc_chrominance_huffman_table, macroblock.cb, 1)); - TRY(encode_ac(ac_chrominance_huffman_table, macroblock.cb)); - - TRY(encode_dc(dc_chrominance_huffman_table, macroblock.cr, 2)); - TRY(encode_ac(ac_chrominance_huffman_table, macroblock.cr)); - - if (mode == Mode::CMYK) { - TRY(encode_dc(dc_luminance_huffman_table, macroblock.k, 3)); - TRY(encode_ac(ac_luminance_huffman_table, macroblock.k)); + for (auto [i, component] : enumerate(to_array({ macroblock.y, macroblock.cb, macroblock.cr, macroblock.k }))) { + if (mode == Mode::CMYK || i < 3) { + TRY(encode_dc(component, i)); + TRY(encode_ac(component, i)); + } } } + m_macroblocks.clear(); - TRY(m_bit_stream.align_to_byte_boundary(0xFF)); - + TRY(find_optimal_huffman_tables()); return {}; } - void set_luminance_quantization_table(QuantizationTable const& table, int quality) + ErrorOr write_huffman_stream() { - set_quantization_table(m_luminance_quantization_table, table, quality); + TRY(write_symbols_to_stream()); + TRY(m_bit_stream.align_to_byte_boundary(0xFF)); + return {}; } - void set_chrominance_quantization_table(QuantizationTable const& table, int quality) + void set_quantization_tables(int quality) { - set_quantization_table(m_chrominance_quantization_table, table, quality); + set_quantization_table(m_luminance_quantization_table, s_default_luminance_quantization_table, quality); + set_quantization_table(m_chrominance_quantization_table, s_default_chrominance_quantization_table, quality); } QuantizationTable const& luminance_quantization_table() const @@ -308,6 +305,17 @@ class JPEGEncodingContext { OutputHuffmanTable ac_chrominance_huffman_table; private: + struct RawBits { + u16 bits {}; + u8 length {}; + }; + struct Symbol { + u8 byte {}; + u8 component_id {}; + bool is_dc {}; + }; + using SymbolOrRawBits = Variant; + static void set_quantization_table(QuantizationTable& destination, QuantizationTable const& source, int quality) { // In order to be compatible with libjpeg-turbo, we use the same coefficients as them. @@ -331,23 +339,36 @@ class JPEGEncodingContext { return m_bit_stream.write_bits(symbol.word, symbol.code_length); } - ErrorOr encode_dc(OutputHuffmanTable const& dc_table, i16 const component[], u8 component_id) + ErrorOr append_symbol(Symbol symbol) + { + auto& stat_table = [&]() -> auto& { + if (symbol.component_id == 0 || symbol.component_id == 3) { + return symbol.is_dc ? m_symbol_stats[0] : m_symbol_stats[1]; + } + return symbol.is_dc ? m_symbol_stats[2] : m_symbol_stats[3]; + }(); + stat_table[symbol.byte] += 1; + TRY(m_symbols_and_bits.try_append(symbol)); + return {}; + } + + ErrorOr encode_dc(i16 const component[], u8 component_id) { // F.1.2.1.3 - Huffman encoding procedures for DC coefficients auto diff = component[0] - m_last_dc_values[component_id]; m_last_dc_values[component_id] = component[0]; auto const size = csize(diff); - TRY(write_symbol(dc_table.from_input_byte(size))); + TRY(append_symbol({ .byte = size, .component_id = component_id, .is_dc = true })); if (diff < 0) diff -= 1; - TRY(m_bit_stream.write_bits(diff, size)); + TRY(m_symbols_and_bits.try_append(RawBits(diff, size))); return {}; } - ErrorOr encode_ac(OutputHuffmanTable const& ac_table, i16 const component[]) + ErrorOr encode_ac(i16 const component[], u8 component_id) { // F.2 - Procedure for sequential encoding of AC coefficients with Huffman coding u32 k {}; @@ -359,7 +380,7 @@ class JPEGEncodingContext { auto coefficient = component[zigzag_map[k]]; if (coefficient == 0) { if (k == 63) { - TRY(write_symbol(ac_table.from_input_byte(0x00))); + TRY(append_symbol({ .byte = 0x00, .component_id = component_id, .is_dc = false })); break; } r += 1; @@ -367,20 +388,20 @@ class JPEGEncodingContext { } while (r > 15) { - TRY(write_symbol(ac_table.from_input_byte(0xF0))); + TRY(append_symbol({ .byte = 0xF0, .component_id = component_id, .is_dc = false })); r -= 16; } { // F.3 - Sequential encoding of a non-zero AC coefficient auto const ssss = csize(coefficient); - auto const rs = (r << 4) + ssss; - TRY(write_symbol(ac_table.from_input_byte(rs))); + u8 const rs = (r << 4) + ssss; + TRY(append_symbol({ .byte = rs, .component_id = component_id, .is_dc = false })); if (coefficient < 0) coefficient -= 1; - TRY(m_bit_stream.write_bits(coefficient, ssss)); + TRY(m_symbols_and_bits.try_append(RawBits(coefficient, ssss))); } r = 0; @@ -388,6 +409,207 @@ class JPEGEncodingContext { return {}; } + ErrorOr write_symbols_to_stream() + { + for (auto const& symbol_or_bits : m_symbols_and_bits) { + if (symbol_or_bits.has()) { + auto symbol = symbol_or_bits.get(); + auto const& huffman_table = [&]() { + if (symbol.component_id == 0 || symbol.component_id == 3) + return symbol.is_dc ? dc_luminance_huffman_table : ac_luminance_huffman_table; + return symbol.is_dc ? dc_chrominance_huffman_table : ac_chrominance_huffman_table; + }(); + TRY(write_symbol(huffman_table.from_input_byte(symbol.byte))); + } else { + auto bits = symbol_or_bits.get(); + TRY(m_bit_stream.write_bits(bits.bits, bits.length)); + } + } + + return {}; + } + + static void find_smallest_frequencies(Array const& frequencies, u16& v1, Optional& v2) + { + // FIXME: A min-heap with a custom comparator should be able to do the trick. + + // "The procedure “Find V1 for least value of FREQ(V1) > 0” always selects the value + // with the largest value of V1 when more than one V1 with the same frequency occurs. + // The reserved code point is then guaranteed to be in the longest code word category." + + u16 index_min {}; + u16 second_index_min {}; + u32 freq_min = NumericLimits::max(); + u32 second_freq_min = NumericLimits::max(); + + for (auto [i, freq] : enumerate(frequencies)) { + if (freq == 0) + continue; + if (freq <= freq_min) { + second_index_min = index_min; + second_freq_min = freq_min; + index_min = i; + freq_min = freq; + } else if (freq <= second_freq_min) { + second_index_min = i; + second_freq_min = freq; + } + } + + v1 = index_min; + if (second_freq_min != NumericLimits::max()) + v2 = second_index_min; + else + v2.clear(); + } + + static Array find_huffman_code_size(Array frequencies) + { + // "Before starting the procedure, the values of FREQ are collected for V = 0 to 255 + // and the FREQ value for V = 256 is set to 1." + frequencies[256] = 1; + + // "the entries in CODESIZE are all set to 0" + Array code_size {}; + + // "the indices in OTHERS are set to –1" + Array others {}; + others.fill(-1); + + // Figure K.1 – Procedure to find Huffman code sizes + while (true) { + u16 v1 {}; + Optional maybe_v2 {}; + find_smallest_frequencies(frequencies, v1, maybe_v2); + if (!maybe_v2.has_value()) + break; + + auto v2 = maybe_v2.value(); + + frequencies[v1] += frequencies[v2]; + frequencies[v2] = 0; + + increment_v1_code_size: + code_size[v1] += 1; + + if (others[v1] != -1) { + v1 = others[v1]; + goto increment_v1_code_size; + } + + others[v1] = v2; + + increment_v2_code_size: + code_size[v2] += 1; + if (others[v2] != -1) { + v2 = others[v2]; + goto increment_v2_code_size; + } + } + + return code_size; + } + + static void adjust_bits(Array& bits) + { + // Figure K.3 – Procedure for limiting code lengths to 16 bits + u16 i = 32; + while (true) { + if (bits[i] > 0) { + auto j = i - 1; + do { + j--; + } while (bits[j] == 0); + + bits[i] = bits[i] - 2; + bits[i - 1] = bits[i - 1] + 1; + bits[j + 1] = bits[j + 1] + 2; + bits[j] = bits[j] - 1; + } else { + i -= 1; + if (i != 16) + continue; + + while (bits[i] == 0) + --i; + bits[i] -= 1; + break; + } + } + } + + static Array count_bits(Array const& code_size) + { + // "The count for each size is contained in the list, BITS. The counts in BITS are zero + // at the start of the procedure." + Array bits {}; + + // Figure K.2 – Procedure to find the number of codes of each size + for (u16 i = 0; i < 257; ++i) { + if (code_size[i] == 0) + continue; + bits[code_size[i]] += 1; + } + adjust_bits(bits); + + return bits; + } + + static Vector sort_input(Array const& code_size) + { + // "Figure K.4 – Sorting of input values according to code size" + Vector huffval {}; + for (u8 i = 1; i <= 32; ++i) { + for (u16 j = 0; j <= 255; ++j) { + if (code_size[j] == i) + huffval.append(j); + } + } + return huffval; + } + + static ErrorOr compute_optimal_table(Array const& distribution) + { + // K.2 A procedure for generating the lists which specify a Huffman code table + + auto code_size = find_huffman_code_size(distribution); + + auto bits = count_bits(code_size); + + // "The input values are sorted according to code size" + auto huffval = sort_input(code_size); + + // "At this point, the list of code lengths (BITS) and the list of values + // (HUFFVAL) can be used to generate the code tables." + + Vector symbols; + u16 code = 0; + u32 symbol_index = 0; + for (auto [encoded_size, number_of_codes] : enumerate(bits)) { + for (u8 i = 0; i < number_of_codes; i++) { + TRY(symbols.try_append({ .input_byte = huffval[symbol_index], .code_length = static_cast(encoded_size), .word = code })); + code++; + symbol_index++; + } + code <<= 1; + } + + return OutputHuffmanTable { move(symbols) }; + } + + ErrorOr find_optimal_huffman_tables() + { + dc_luminance_huffman_table = TRY(compute_optimal_table(m_symbol_stats[0])); + dc_luminance_huffman_table.id = (0 << 4) | 0; + ac_luminance_huffman_table = TRY(compute_optimal_table(m_symbol_stats[1])); + ac_luminance_huffman_table.id = (1 << 4) | 0; + dc_chrominance_huffman_table = TRY(compute_optimal_table(m_symbol_stats[2])); + dc_chrominance_huffman_table.id = (0 << 4) | 1; + ac_chrominance_huffman_table = TRY(compute_optimal_table(m_symbol_stats[3])); + ac_chrominance_huffman_table.id = (1 << 4) | 1; + return {}; + } + static u8 csize(i16 coefficient) { VERIFY(coefficient >= -2047 && coefficient <= 2047); @@ -404,6 +626,9 @@ class JPEGEncodingContext { Vector m_macroblocks {}; Array m_last_dc_values {}; + Array, 4> m_symbol_stats {}; + Vector m_symbols_and_bits {}; + JPEGBigEndianOutputBitStream m_bit_stream; }; @@ -601,17 +826,8 @@ ErrorOr add_scan_header(Stream& stream, Mode mode) return {}; } -ErrorOr add_headers(Stream& stream, JPEGEncodingContext& context, JPEGWriter::Options const& options, IntSize size, Mode mode) +ErrorOr add_headers(Stream& stream, JPEGEncodingContext const& context, JPEGWriter::Options const& options, IntSize size, Mode mode) { - context.set_luminance_quantization_table(s_default_luminance_quantization_table, options.quality); - context.set_chrominance_quantization_table(s_default_chrominance_quantization_table, options.quality); - - context.dc_luminance_huffman_table = s_default_dc_luminance_huffman_table; - context.dc_chrominance_huffman_table = s_default_dc_chrominance_huffman_table; - - context.ac_luminance_huffman_table = s_default_ac_luminance_huffman_table; - context.ac_chrominance_huffman_table = s_default_ac_chrominance_huffman_table; - TRY(add_start_of_image(stream)); if (options.icc_data.has_value()) @@ -633,11 +849,14 @@ ErrorOr add_headers(Stream& stream, JPEGEncodingContext& context, JPEGWrit return {}; } -ErrorOr add_image(Stream& stream, JPEGEncodingContext& context, Mode mode) +ErrorOr add_image(Stream& stream, JPEGEncodingContext& context, JPEGEncoderOptions const& options, IntSize size, Mode mode) { + context.set_quantization_tables(options.quality); context.convert_to_ycbcr(mode); context.fdct_and_quantization(mode); - TRY(context.write_huffman_stream(mode)); + TRY(context.huffman_encode_macroblocks(mode)); + TRY(add_headers(stream, context, options, size, mode)); + TRY(context.write_huffman_stream()); TRY(add_end_of_image(stream)); return {}; } @@ -647,18 +866,16 @@ ErrorOr add_image(Stream& stream, JPEGEncodingContext& context, Mode mode) ErrorOr JPEGWriter::encode(Stream& stream, Bitmap const& bitmap, Options const& options) { JPEGEncodingContext context { JPEGBigEndianOutputBitStream { stream } }; - TRY(add_headers(stream, context, options, bitmap.size(), Mode::RGB)); TRY(context.initialize_mcu(bitmap)); - TRY(add_image(stream, context, Mode::RGB)); + TRY(add_image(stream, context, options, bitmap.size(), Mode::RGB)); return {}; } ErrorOr JPEGWriter::encode(Stream& stream, CMYKBitmap const& bitmap, Options const& options) { JPEGEncodingContext context { JPEGBigEndianOutputBitStream { stream } }; - TRY(add_headers(stream, context, options, bitmap.size(), Mode::CMYK)); TRY(context.initialize_mcu(bitmap)); - TRY(add_image(stream, context, Mode::CMYK)); + TRY(add_image(stream, context, options, bitmap.size(), Mode::CMYK)); return {}; }