dlib C++ Library - save_jxl.cpp

    // Copyright (C) 2024  Davis E. King (davis@dlib.net), Adrià Arrufat
    // License: Boost Software License   See LICENSE.txt for the full license.
    #ifndef DLIB_JXL_SAVER_CPp_
    #define DLIB_JXL_SAVER_CPp_

    // only do anything with this file if DLIB_JXL_SUPPORT is defined
    #ifdef DLIB_JXL_SUPPORT

    #include "save_jxl.h"
    #include "image_saver.h"
    #include <sstream>
    #include <jxl/encode_cxx.h>
    #include <jxl/resizable_parallel_runner_cxx.h>

    namespace dlib {

    // ----------------------------------------------------------------------------------------

        namespace impl
        {
            void impl_save_jxl (
                const std::string& filename,
                const uint8_t* pixels,
                const uint32_t width,
                const uint32_t height,
                const uint32_t num_channels,
                const float quality
            )
            {
                std::ofstream fout(filename, std::ios::binary);
                if (!fout.good())
                {
                    throw image_save_error("Unable to open " + filename + " for writing.");
                }

                auto enc = JxlEncoderMake(nullptr);
                auto runner = JxlResizableParallelRunnerMake(nullptr);
                JxlResizableParallelRunnerSetThreads(runner.get(), JxlResizableParallelRunnerSuggestThreads(width, height));

                if (JXL_ENC_SUCCESS != JxlEncoderSetParallelRunner(enc.get(), JxlResizableParallelRunner, runner.get()))
                {
                    throw image_save_error("jxl_saver: JxlResizableParallelRunner failed");
                }

                JxlPixelFormat pixel_format{
                    .num_channels = num_channels,
                    .data_type = JXL_TYPE_UINT8,
                    .endianness = JXL_NATIVE_ENDIAN,
                    .align = 0
                };
                JxlBasicInfo basic_info;
                JxlEncoderInitBasicInfo(&basic_info);
                basic_info.xsize = width;
                basic_info.ysize = height;
                basic_info.bits_per_sample = 8;
                basic_info.uses_original_profile = quality == 100;
                switch (num_channels)
                {
                case 1:
                    basic_info.num_color_channels = 1;
                    basic_info.num_extra_channels = 0;
                    basic_info.alpha_bits = 0;
                    basic_info.alpha_exponent_bits = 0;
                    break;
                case 3:
                    basic_info.num_color_channels = 3;
                    basic_info.num_extra_channels = 0;
                    basic_info.alpha_bits = 0;
                    basic_info.alpha_exponent_bits = 0;
                    break;
                case 4:
                    basic_info.num_color_channels = 3;
                    basic_info.num_extra_channels = 1;
                    basic_info.alpha_bits = basic_info.bits_per_sample;
                    basic_info.alpha_exponent_bits = 0;
                    break;
                default:
                    throw ("jxl_saver: unsupported number of channels");
                }

                if (JXL_ENC_SUCCESS != JxlEncoderSetBasicInfo(enc.get(), &basic_info))
                {
                    throw image_save_error("jxl_saver: JxlEncoderSetBasicInfo failed");
                }

                JxlColorEncoding color_encoding = {};
                JxlColorEncodingSetToSRGB(&color_encoding, /* is_gray = */ num_channels < 3);
                if (JXL_ENC_SUCCESS != JxlEncoderSetColorEncoding(enc.get(), &color_encoding))
                {
                    throw image_save_error("jxl_saver: JxlEncoderSetColorEncoding failed");
                }

                JxlEncoderFrameSettings* frame_settings = JxlEncoderFrameSettingsCreate(enc.get(), nullptr);
                JxlEncoderFrameSettingsSetOption(frame_settings, JXL_ENC_FRAME_SETTING_DECODING_SPEED, 0);

                const float distance = JxlEncoderDistanceFromQuality(quality);
                if (JXL_ENC_SUCCESS != JxlEncoderSetFrameDistance(frame_settings, distance))
                {
                    throw image_save_error("jxl_saver: JxlEncoderSetFrameDistance failed");
                }
                if (basic_info.alpha_bits > 0)
                {
                    if (JXL_ENC_SUCCESS != JxlEncoderSetExtraChannelDistance(frame_settings, 0, distance))
                    {
                        throw image_save_error("jxl_saver: JxlEncoderSetExtraChannelDistance failed");
                    }
                }

                // explictly enable lossless mode
                if (distance == 0)
                {
                    if (JXL_ENC_SUCCESS != JxlEncoderSetFrameLossless(frame_settings, JXL_TRUE))
                    {
                        throw image_save_error("jxl_saver: JxlEncoderSetFrameLossless failed");
                    }
                }

                void* pixels_data = reinterpret_cast<void*>(const_cast<uint8_t*>(pixels));
                const size_t pixels_size = width * height * num_channels;
                if (JXL_ENC_SUCCESS != JxlEncoderAddImageFrame(frame_settings, &pixel_format, pixels_data, pixels_size))
                {
                    throw image_save_error("jxl_saver: JxlEncoderAddImageFrame failed");
                }
                JxlEncoderCloseInput(enc.get());

                std::vector<uint8_t> compressed;
                compressed.resize(64);
                uint8_t* next_out = compressed.data();
                size_t avail_out = compressed.size() - (next_out - compressed.data());
                JxlEncoderStatus process_result = JXL_ENC_NEED_MORE_OUTPUT;
                while (process_result == JXL_ENC_NEED_MORE_OUTPUT)
                {
                    process_result = JxlEncoderProcessOutput(enc.get(), &next_out, &avail_out);
                    if (process_result == JXL_ENC_NEED_MORE_OUTPUT)
                    {
                        size_t offset = next_out - compressed.data();
                        compressed.resize(compressed.size() * 2);
                        next_out = compressed.data() + offset;
                        avail_out = compressed.size() - offset;
                    }
                }
                compressed.resize(next_out - compressed.data());
                if (JXL_ENC_SUCCESS != process_result)
                {
                    throw image_save_error("jxl_saver: JxlEncoderProcessOutput failed");
                }
                fout.write(reinterpret_cast<char*>(compressed.data()), compressed.size());
                if (!fout.good())
                {
                    throw image_save_error("Error while writing JPEG XL image to " + filename + ".");
                }
            }
        }

    // ----------------------------------------------------------------------------------------

    }

    #endif // DLIB_JXL_SUPPORT

    #endif // DLIB_JXL_SAVER_CPp_