png.hpp

 
#ifndef UFO_IO_PNG_HPP
#define UFO_IO_PNG_HPP
 
// UFO
#include <ufo/io/file_handler.hpp>
#include <ufo/io/image_properties.hpp>
#include <ufo/vision/color.hpp>
#include <ufo/vision/image.hpp>
 
// STL
#include <algorithm>
#include <cstdio>
#include <filesystem>
#include <print>
#include <type_traits>
 
namespace ufo
{
namespace detail
{
bool readPNG(std::FILE* fp, void* image, std::size_t length, int fmt);
 
bool writePNG(std::FILE* fp, void const* image, std::size_t length, std::uint32_t width,
              std::uint32_t height, int color_type, int bit_depth);
}  // namespace detail
 
[[nodiscard]] ImageProperties imagePropertiesPNG(std::filesystem::path const& file);
 
template <ColorType CT, class T, bool Alpha, bool Weight>
bool readPNG(std::filesystem::path const& file, Image<Color<CT, T, Alpha, Weight>>& image)
{
    FileHandler fp(file, "rb");
 
    if (!fp) {
        std::println(stderr, "[UFO | Read PNG] Failed to open file: {}", file.string());
        return false;
    }
 
    auto fun = [&image](std::FILE* fp, auto& tmp, int fmt) {
        using C = typename std::remove_cvref_t<decltype(tmp)>::value_type;
        if (detail::readPNG(fp, tmp.data(), sizeof(C) * tmp.size(), fmt)) {
            if constexpr (has_alpha_v<C>) {
                for (auto& p : tmp) {
                    if (0 == alpha(p)) {
                        if constexpr (ColorType::GRAY == color_type_v<C>) {
                            p.gray = 0;
                        } else {
                            p.red   = 0;
                            p.green = 0;
                            p.blue  = 0;
                        }
                    }
                }
            }
            image = convert<Color<CT, T, Alpha, Weight>>(tmp);
            return true;
        }
        return false;
    };
 
    auto const prop = imagePropertiesPNG(file);
    if (prop.alpha && prop.grayscale && 8 >= prop.bit_depth) {
        int                                                      fmt = 16;  // SPNG_FMT_GA8
        Image<Color<ColorType::GRAY, std::uint8_t, true, false>> tmp(prop.height, prop.width);
        return fun(fp.get(), tmp, fmt);
    } else if (prop.grayscale && 16 == prop.bit_depth) {
        int                                                       fmt = 32;  // SPNG_FMT_GA16
        Image<Color<ColorType::GRAY, std::uint16_t, true, false>> tmp(prop.height,
                                                                      prop.width);
        return fun(fp.get(), tmp, fmt);
    } else if (prop.alpha && !prop.grayscale && 8 == prop.bit_depth) {
        int                                                     fmt = 1;  // SPNG_FMT_RGBA8
        Image<Color<ColorType::RGB, std::uint8_t, true, false>> tmp(prop.height, prop.width);
        return fun(fp.get(), tmp, fmt);
    } else if (!prop.grayscale && 16 == prop.bit_depth) {
        int                                                      fmt = 2;  // SPNG_FMT_RGBA16
        Image<Color<ColorType::RGB, std::uint16_t, true, false>> tmp(prop.height, prop.width);
        return fun(fp.get(), tmp, fmt);
    } else if (!prop.alpha && prop.grayscale && 8 >= prop.bit_depth) {
        int                                                       fmt = 64;  // SPNG_FMT_G8
        Image<Color<ColorType::GRAY, std::uint8_t, false, false>> tmp(prop.height,
                                                                      prop.width);
        return fun(fp.get(), tmp, fmt);
    } else if (!prop.alpha && !prop.grayscale && 8 == prop.bit_depth) {
        int                                                      fmt = 4;  // SPNG_FMT_RGB8
        Image<Color<ColorType::RGB, std::uint8_t, false, false>> tmp(prop.height, prop.width);
        return fun(fp.get(), tmp, fmt);
    } else {
        std::println(stderr, "[UFO | Read PNG] Unknown PNG type");
        return false;
    }
}
 
template <ColorType CT, class T, bool Alpha, bool Weight>
bool writePNG(std::filesystem::path const&              file,
              Image<Color<CT, T, Alpha, Weight>> const& image)
{
    FileHandler fp(file, "wb");
 
    if (!fp) {
        std::println(stderr, "[UFO | Write PNG] Failed to create file: {}", file.string());
        return false;
    }
 
    std::uint32_t const width     = image.cols();
    std::uint32_t const height    = image.rows();
    int const           bit_depth = sizeof(std::uint8_t) < sizeof(T) ? 16 : 8;
 
    // SPNG colour type and channel count derived entirely from template parameters.
    constexpr int color_type = (ColorType::GRAY == CT) ? (Alpha ? 4 : 0)   // GA or G
                                                       : (Alpha ? 6 : 2);  // RGBA or RGB
    constexpr int channels   = (ColorType::GRAY == CT) ? (Alpha ? 2 : 1) : (Alpha ? 4 : 3);
 
    std::size_t const length =
        static_cast<std::size_t>(width) * height * (bit_depth / 8) * channels;
 
    // Canonical output types: strip Weight, normalise scalar width.
    constexpr ColorType CT2 = (ColorType::GRAY == CT) ? ColorType::GRAY : ColorType::RGB;
    using T2 =
        std::conditional_t<sizeof(std::uint8_t) < sizeof(T), std::uint16_t, std::uint8_t>;
 
    if constexpr (CT == CT2 && std::is_same_v<T, T2> && !Weight) {
        return detail::writePNG(fp.get(), image.data(), length, width, height, color_type,
                                bit_depth);
    } else {
        Image<Color<CT2, T2, Alpha, false>> data(image.rows(), image.cols());
        std::ranges::transform(image, data.begin(), [](auto x) {
            return convert<Color<CT2, T2, Alpha, false>>(x);
        });
        return detail::writePNG(fp.get(), data.data(), length, width, height, color_type,
                                bit_depth);
    }
}
}  // namespace ufo
 
#endif  // UFO_IO_PNG_HPP