o3de/Gems/Atom/Utils/Code/Source/PngFile.cpp

/*  
 * Copyright (c) Contributors to the Open 3D Engine Project.
 * For complete copyright and license terms please see the LICENSE at the root of this distribution.
 *
 * SPDX-License-Identifier: Apache-2.0 OR MIT
 *
 */

#include <Atom/Utils/PngFile.h>
#include <png.h>

namespace AZ
{
    namespace
    {
        void PngImage_user_error_fn(png_structp png_ptr, png_const_charp error_msg)
        {
            PngImage::ErrorHandler* errorHandler = reinterpret_cast<PngImage::ErrorHandler*>(png_get_error_ptr(png_ptr));
            (*errorHandler)(error_msg);
        }

        void PngImage_user_warning_fn(png_structp /*png_ptr*/, png_const_charp warning_msg)
        {
            AZ_Warning("PngImage", false, "%s", warning_msg);
        }
    }

    PngImage PngImage::Create(const RHI::Size& size, RHI::Format format, AZStd::array_view<uint8_t> data)
    {
        return Create(size, format, AZStd::vector<uint8_t>{data.begin(), data.end()});
    }

    PngImage PngImage::Create(const RHI::Size& size, RHI::Format format, AZStd::vector<uint8_t>&& data)
    {
        PngImage image;

        if (RHI::Format::R8G8B8A8_UNORM == format)
        {
            if (size.m_width * size.m_height * 4 == data.size())
            {
                image.m_width = size.m_width;
                image.m_height = size.m_height;
                image.m_bitDepth = 8;
                image.m_colorType = PNG_COLOR_TYPE_RGB_ALPHA;
                image.m_bufferFormat = PngImage::Format::RGBA;
                image.m_buffer = data;
            }
            else
            {
                AZ_Assert(false, "Invalid arguments. Buffer size does not match the image dimensions.");
            }
        }

        return image;
    }
    
    PngImage PngImage::Load(const char* path, LoadSettings loadSettings)
    {
        if (!loadSettings.m_errorHandler)
        {
            loadSettings.m_errorHandler = [path](const char* message) { DefaultErrorHandler(AZStd::string::format("Could not load file '%s'. %s", path, message).c_str()); };
        }
        
        // For documentation of this code, see http://www.libpng.org/pub/png/libpng-1.4.0-manual.pdf chapter 3

        FILE* fp = NULL;
        if (fopen_s(&fp, path, "rb") || !fp)
        {
            loadSettings.m_errorHandler("Failed to open file.");
            return {};
        }

        png_byte header[HeaderSize] = {};

        if (fread(header, 1, HeaderSize, fp) != HeaderSize)
        {
            fclose(fp);
            loadSettings.m_errorHandler("Invalid header.");
            return {};
        }

        bool isPng = !png_sig_cmp(header, 0, HeaderSize);
        if (!isPng)
        {
            fclose(fp);
            loadSettings.m_errorHandler("Invalid header.");
            return {};
        }

        png_voidp user_error_ptr = &loadSettings.m_errorHandler;
        png_error_ptr user_error_fn = PngImage_user_error_fn;
        png_error_ptr user_warning_fn = PngImage_user_warning_fn;

        png_structp png_ptr = png_create_read_struct(PNG_LIBPNG_VER_STRING, user_error_ptr, user_error_fn, user_warning_fn);
        if (!png_ptr)
        {
            fclose(fp);
            loadSettings.m_errorHandler("png_create_read_struct failed.");
            return {};
        }

        png_infop info_ptr = png_create_info_struct(png_ptr);
        if (!info_ptr)
        {
            png_destroy_read_struct(&png_ptr, (png_infopp)NULL, (png_infopp)NULL);
            fclose(fp);
            loadSettings.m_errorHandler("png_create_info_struct failed.");
            return {};
        }

        png_infop end_info = png_create_info_struct(png_ptr);
        if (!end_info)
        {
            png_destroy_read_struct(&png_ptr, &info_ptr, (png_infopp)NULL);
            fclose(fp);
            loadSettings.m_errorHandler("png_create_info_struct failed.");
            return {};
        }

#pragma warning(push)
#pragma warning(disable: 4611) // Disables "interaction between '_setjmp' and C++ object destruction is non-portable". See https://docs.microsoft.com/en-us/cpp/preprocessor/warning?view=msvc-160
        if (setjmp(png_jmpbuf(png_ptr)))
        {
            png_destroy_read_struct(&png_ptr, &info_ptr, &end_info);
            fclose(fp);
            // We don't report an error message here because the user_error_fn should have done that already.
            return {};
        }
#pragma warning(pop)

        png_init_io(png_ptr, fp);

        png_set_sig_bytes(png_ptr, HeaderSize);

        png_set_keep_unknown_chunks(png_ptr, PNG_HANDLE_CHUNK_NEVER, NULL, 0);

        // To keep things simple for now we limit all images to RGB and RGBA, 8 bits per channel
        int png_transforms = PNG_TRANSFORM_PACKING | // Expand 1, 2 and 4-bit samples to bytes
                             PNG_TRANSFORM_STRIP_16 | // Reduce 16 bit samples to 8 bits
                             PNG_TRANSFORM_GRAY_TO_RGB;

        if (loadSettings.m_stripAlpha)
        {
            png_transforms |= PNG_TRANSFORM_STRIP_ALPHA;
        }

        png_read_png(png_ptr, info_ptr, png_transforms, NULL);
        
        // Note that libpng will allocate row_pointers for us. If we want to manage the memory ourselves, we need to call png_set_rows.
        // In that case we would have to use the low level interface: png_read_info, png_read_image, and png_read_end.
        png_bytep* row_pointers = png_get_rows(png_ptr, info_ptr);

        PngImage pngImage;

        png_get_IHDR(png_ptr, info_ptr, &pngImage.m_width, &pngImage.m_height, &pngImage.m_bitDepth, &pngImage.m_colorType, NULL, NULL, NULL);
        
        uint32_t bytesPerPixel = 0;

        switch (pngImage.m_colorType)
        {
        case PNG_COLOR_TYPE_RGB:
            pngImage.m_bufferFormat = PngImage::Format::RGB;
            bytesPerPixel = 3;
            break;
        case PNG_COLOR_TYPE_RGBA:
            pngImage.m_bufferFormat = PngImage::Format::RGBA;
            bytesPerPixel = 4;
            break;
        default:
            AZ_Assert(false, "The png transforms should have ensured a pixel format of RGB or RGBA, 8 bits per channel");
            png_destroy_read_struct(&png_ptr, &info_ptr, (png_infopp)NULL);
            fclose(fp);
            loadSettings.m_errorHandler("Unsupported pixel format.");
            return {};
        }

        // In the future we could use the low-level interface to avoid copying the image (and provide progress callbacks)
        pngImage.m_buffer.set_capacity(pngImage.m_width * pngImage.m_height * bytesPerPixel);
        for (uint32_t rowIndex = 0; rowIndex < pngImage.m_height; ++rowIndex)
        {
            png_bytep row = row_pointers[rowIndex];
            pngImage.m_buffer.insert(pngImage.m_buffer.end(), row, row + (pngImage.m_width * bytesPerPixel));
        }
        
        png_destroy_read_struct(&png_ptr, &info_ptr, &end_info);
        fclose(fp);
        return pngImage;
    }

    bool PngImage::Save(const char* path, SaveSettings saveSettings)
    {
        if (!saveSettings.m_errorHandler)
        {
            saveSettings.m_errorHandler = [path](const char* message) { DefaultErrorHandler(AZStd::string::format("Could not save file '%s'. %s", path, message).c_str()); };
        }
        
        if (!IsValid())
        {
            saveSettings.m_errorHandler("This PngImage is invalid.");
            return false;
        }

        // For documentation of this code, see http://www.libpng.org/pub/png/libpng-1.4.0-manual.pdf chapter 4

        FILE* fp = NULL;
        if (fopen_s(&fp, path, "wb") || !fp)
        {
            saveSettings.m_errorHandler("Failed to open file.");
            return false;
        }
        
        png_voidp user_error_ptr = &saveSettings.m_errorHandler;
        png_error_ptr user_error_fn = PngImage_user_error_fn;
        png_error_ptr user_warning_fn = PngImage_user_warning_fn;

        png_structp png_ptr = png_create_write_struct(PNG_LIBPNG_VER_STRING, user_error_ptr, user_error_fn, user_warning_fn);
        if (!png_ptr)
        {
            fclose(fp);
            saveSettings.m_errorHandler("png_create_write_struct failed.");
            return false;
        }

        png_infop info_ptr = png_create_info_struct(png_ptr);
        if (!info_ptr)
        {
            png_destroy_write_struct(&png_ptr, (png_infopp)NULL);
            fclose(fp);
            saveSettings.m_errorHandler("png_destroy_write_struct failed.");
            return false;
        }

#pragma warning(push)
#pragma warning(disable: 4611) // Disables "interaction between '_setjmp' and C++ object destruction is non-portable". See https://docs.microsoft.com/en-us/cpp/preprocessor/warning?view=msvc-160
        if (setjmp(png_jmpbuf(png_ptr)))
        {
            png_destroy_write_struct(&png_ptr, &info_ptr);
            fclose(fp);
            // We don't report an error message here because the user_error_fn should have done that already.
            return false;
        }
#pragma warning(pop)

        png_init_io(png_ptr, fp);

        png_set_IHDR(png_ptr, info_ptr, m_width, m_height, m_bitDepth, m_colorType, PNG_INTERLACE_NONE, PNG_COMPRESSION_TYPE_DEFAULT, PNG_FILTER_TYPE_DEFAULT);
        
        png_set_compression_level(png_ptr, saveSettings.m_compressionLevel);

        const uint32_t bytesPerPixel = (m_bufferFormat == PngImage::Format::RGB) ? 3 : 4;

        AZStd::vector<uint8_t*> rows;
        rows.reserve(m_height);
        for (uint32_t i = 0; i < m_height; ++i)
        {
            rows.push_back(m_buffer.begin() + m_width * bytesPerPixel * i);
        }

        png_set_rows(png_ptr, info_ptr, rows.begin());

        int transforms = PNG_TRANSFORM_IDENTITY;
        if (saveSettings.m_stripAlpha && m_bufferFormat == PngImage::Format::RGBA)
        {
            transforms |= PNG_TRANSFORM_STRIP_FILLER_AFTER;
        }

        png_write_png(png_ptr, info_ptr, PNG_TRANSFORM_IDENTITY, NULL);

        png_destroy_write_struct(&png_ptr, &info_ptr);
        
        fclose(fp);

        return true;
    }
    
    void PngImage::DefaultErrorHandler(const char* message)
    {
        AZ_Error("PngImage", false, "%s", message);
    }

    bool PngImage::IsValid() const
    {
        return
            !m_buffer.empty() &&
            m_width > 0 &&
            m_height > 0 &&
            m_bitDepth > 0;
    }

    AZStd::vector<uint8_t>&& PngImage::TakeBuffer()
    {
        return AZStd::move(m_buffer);
    }
        
}// namespace AZ