dlib C++ Library - jxl_loader.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_LOADER_CPp_
    #define DLIB_JXL_LOADER_CPp_

    // only do anything with this file if DLIB_JXL_SUPPORT is defined
    #ifdef DLIB_JXL_SUPPORT
    #include "jxl_loader.h"
    #include <jxl/decode_cxx.h>
    #include <jxl/resizable_parallel_runner_cxx.h>
    #include <fstream>

    namespace dlib
    {

        static std::vector<unsigned char> load_contents(const std::string& filename)
        {
            std::ifstream stream(filename, std::ios::binary);
            stream.exceptions(std::ifstream::failbit | std::ifstream::badbit | std::ifstream::eofbit);
            std::vector<unsigned char> buffer;
            vectorstream temp(buffer);
            temp << stream.rdbuf();
            return buffer;
        }

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

        jxl_loader::
        jxl_loader(const char* filename) : height(0), width(0)
        {
            data = load_contents(filename);
            get_info();
        }

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

        jxl_loader::
        jxl_loader(const std::string& filename) : height(0), width(0)
        {
            data = load_contents(filename);
            get_info();
        }

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

        jxl_loader::
        jxl_loader(const dlib::file& f) : height(0), width(0)
        {
            data = load_contents(f.full_name());
            get_info();
        }

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

        jxl_loader::
        jxl_loader(const unsigned char* imgbuffer, size_t imgbuffersize) : height(0), width(0)
        {
            data.resize(imgbuffersize);
            memcpy(data.data(), imgbuffer, imgbuffersize);
            get_info();
        }

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

        bool jxl_loader::is_gray() const { return depth == 1; }
        bool jxl_loader::is_graya() const { return depth == 2; };
        bool jxl_loader::is_rgb() const { return depth == 3; }
        bool jxl_loader::is_rgba() const { return depth == 4; }
        unsigned int jxl_loader::bit_depth() const { return bits_per_sample; };
        long jxl_loader::nr() const { return static_cast<long>(height); };
        long jxl_loader::nc() const { return static_cast<long>(width); };

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

        void jxl_loader::get_info()
        {
            JxlSignature signature = JxlSignatureCheck(data.data(), data.size());
            if (signature != JXL_SIG_CODESTREAM && signature != JXL_SIG_CONTAINER)
            {
                throw image_load_error("jxl_loader: JxlSignatureCheck failed");
            }

            auto dec = JxlDecoderMake(nullptr);
            if (JXL_DEC_SUCCESS != JxlDecoderSubscribeEvents(dec.get(), JXL_DEC_BASIC_INFO))
            {
                throw image_load_error("jxl_loader: JxlDecoderSubscribeEvents failed");
            }

            JxlDecoderSetInput(dec.get(), data.data(), data.size());
            JxlDecoderCloseInput(dec.get());
            if (JXL_DEC_BASIC_INFO != JxlDecoderProcessInput(dec.get())) {
                throw image_load_error("jxl_loader: JxlDecoderProcessInput failed");
            }

            JxlBasicInfo basic_info;
            if (JXL_DEC_SUCCESS != JxlDecoderGetBasicInfo(dec.get(), &basic_info))
            {
                throw image_load_error("jxl_loader: JxlDecoderGetBasicInfo failed");
            }
            width = basic_info.xsize;
            height = basic_info.ysize;
            depth = basic_info.num_color_channels + basic_info.num_extra_channels;
            bits_per_sample = basic_info.bits_per_sample;
        }
    // ----------------------------------------------------------------------------------------

        void jxl_loader::decode(unsigned char* out, const size_t out_size) const
        {
            auto runner = JxlResizableParallelRunnerMake(nullptr);
            auto dec = JxlDecoderMake(nullptr);
            if (JXL_DEC_SUCCESS != JxlDecoderSubscribeEvents(dec.get(), JXL_DEC_FULL_IMAGE))
            {
                throw image_load_error("jxl_loader: JxlDecoderSubscribeEvents failed");
            }

            if (JXL_DEC_SUCCESS != JxlDecoderSetParallelRunner(dec.get(), JxlResizableParallelRunner, runner.get()))
            {
                throw image_load_error("jxl_loader: JxlDecoderSetParallelRunner failed");
            }

            if (JXL_DEC_SUCCESS != JxlDecoderSetInput(dec.get(), data.data(), data.size()))
            {
                throw image_load_error("jxl_loader: JxlDecoderSetInput failed");
            }
            JxlDecoderCloseInput(dec.get());

            JxlPixelFormat format = {
                .num_channels = depth,
                .data_type = JXL_TYPE_UINT8,
                .endianness = JXL_NATIVE_ENDIAN,
                .align=0
            };
            for (;;)
            {
                JxlDecoderStatus status = JxlDecoderProcessInput(dec.get());
                if (status == JXL_DEC_ERROR)
                {
                    throw image_load_error("jxl_loader: JxlDecoderProcessInput failed");
                }
                else if (status == JXL_DEC_NEED_MORE_INPUT)
                {
                    throw image_load_error("jxl_loader: Error, expected more input");
                }
                else if (status == JXL_DEC_NEED_IMAGE_OUT_BUFFER)
                {
                    JxlResizableParallelRunnerSetThreads(runner.get(), JxlResizableParallelRunnerSuggestThreads(width, height));
                    size_t buffer_size;
                    if (JXL_DEC_SUCCESS != JxlDecoderImageOutBufferSize(dec.get(), &format, &buffer_size))
                    {
                        throw image_load_error("jxl_loader: JxlDecoderImageOutBufferSize failed");
                    }
                    if (buffer_size != width * height * depth)
                    {
                        throw image_load_error("jxl_loader: invalid output buffer size");
                    }
                    if (JXL_DEC_SUCCESS != JxlDecoderSetImageOutBuffer(dec.get(), &format, out, out_size))
                    {
                        throw image_load_error("jxl_loader: JxlDecoderSetImageOutBuffer failed");
                    }
                }
                else if (status == JXL_DEC_FULL_IMAGE)
                {
                    // If the image is an animation, more full frames may be decoded.
                    // This loader only decodes the first one.
                    return;
                }
                else if (status == JXL_DEC_SUCCESS)
                {
                    return;
                }
                else
                {
                    throw image_load_error("jxl_loder: Unknown decoder status");
                }
            }
        }
    }

    #endif // DLIB_JXL_SUPPORT
    #endif // DLIB_JXL_LOADER_CPp_