io/image.cpp
| Line | Branch | Exec | Source |
|---|---|---|---|
| 1 | // SPDX-License-Identifier: BSD-3-Clause | ||
| 2 | /// @file image.cpp | ||
| 3 | /// @brief Image loading, EXR writing, sRGB conversion and mip chains. | ||
| 4 | #include "ernie/io/image.h" | ||
| 5 | |||
| 6 | #include "ernie/util/profiler.h" | ||
| 7 | #include "ernie/util/text.h" | ||
| 8 | |||
| 9 | #include <ImfChannelList.h> | ||
| 10 | #include <ImfFrameBuffer.h> | ||
| 11 | #include <ImfHeader.h> | ||
| 12 | #include <ImfOutputFile.h> | ||
| 13 | #include <ImfRgbaFile.h> | ||
| 14 | #include <stb_image.h> | ||
| 15 | |||
| 16 | #include <algorithm> | ||
| 17 | #include <array> | ||
| 18 | #include <cmath> | ||
| 19 | #include <cstddef> | ||
| 20 | #include <cstdlib> | ||
| 21 | #include <stdexcept> | ||
| 22 | #include <string> | ||
| 23 | #include <vector> | ||
| 24 | |||
| 25 | namespace ernie { | ||
| 26 | |||
| 27 | namespace { | ||
| 28 | |||
| 29 | 9 | Image LoadExr(const std::filesystem::path& path) { | |
| 30 | 9 | try { | |
| 31 | // RGBA as half floats, whatever the file holds: missing alpha | ||
| 32 | // reads as 1, luminance images as grey, and every compression | ||
| 33 | // (RLE, ZIP, PIZ, PXR24, B44, DWAA, DWAB, HTJ2K) decodes, | ||
| 34 | // scanline or tiled. | ||
| 35 | 10 | Imf::RgbaInputFile file(path.string().c_str()); | |
| 36 | 8 | const Imath::Box2i window = file.dataWindow(); | |
| 37 | 8 | const int width = window.max.x - window.min.x + 1; | |
| 38 | 8 | const int height = window.max.y - window.min.y + 1; | |
| 39 | 8 | std::vector<Imf::Rgba> texels(static_cast<std::size_t>(width) * | |
| 40 | 8 | static_cast<std::size_t>(height)); | |
| 41 | // Strides that put the window's first pixel at the buffer's | ||
| 42 | // start. | ||
| 43 | 8 | file.setFrameBuffer(texels.data() - window.min.x - | |
| 44 | 8 | static_cast<std::ptrdiff_t>(window.min.y) * width, | |
| 45 | 1, static_cast<std::size_t>(width)); | ||
| 46 | 8 | file.readPixels(window.min.y, window.max.y); | |
| 47 | 8 | Image image(static_cast<std::uint32_t>(width), static_cast<std::uint32_t>(height)); | |
| 48 | 12583 | for (std::size_t i = 0; i < image.pixels.size(); ++i) { | |
| 49 | 12575 | const Imf::Rgba& t = texels[i]; | |
| 50 | 12575 | image.pixels[i] = glm::vec4(static_cast<float>(t.r), static_cast<float>(t.g), | |
| 51 | static_cast<float>(t.b), static_cast<float>(t.a)); | ||
| 52 | } | ||
| 53 | 24 | return image; | |
| 54 | 9 | } catch (const std::exception& error) { | |
| 55 | 4 | throw std::runtime_error("Cannot read " + path.string() + ": " + error.what()); | |
| 56 | 1 | } | |
| 57 | } | ||
| 58 | |||
| 59 | template <typename T> | ||
| 60 | 7 | Image Decode(T* data, int width, int height, float scale, ImageEncoding encoding) { | |
| 61 | 7 | Image image(static_cast<std::uint32_t>(width), static_cast<std::uint32_t>(height)); | |
| 62 | 1908 | for (std::size_t i = 0; i < image.pixels.size(); ++i) { | |
| 63 | glm::vec4 value; | ||
| 64 | 9505 | for (int c = 0; c < 4; ++c) { | |
| 65 | 7604 | value[c] = static_cast<float>(data[4 * i + static_cast<std::size_t>(c)]) * scale; | |
| 66 | } | ||
| 67 | 1901 | if (encoding == ImageEncoding::Srgb) { | |
| 68 | 853 | value = glm::vec4(SrgbToLinear(value.r), SrgbToLinear(value.g), SrgbToLinear(value.b), | |
| 69 | value.a); | ||
| 70 | } | ||
| 71 | 1901 | image.pixels[i] = value; | |
| 72 | } | ||
| 73 | 7 | stbi_image_free(data); | |
| 74 | 7 | return image; | |
| 75 | ✗ | } | |
| 76 | |||
| 77 | 10 | Image LoadStb(const std::filesystem::path& path, ImageEncoding encoding) { | |
| 78 | 10 | const std::string name = path.string(); | |
| 79 | 10 | int width = 0; | |
| 80 | 10 | int height = 0; | |
| 81 | 10 | int channels = 0; | |
| 82 | 10 | if (stbi_is_16_bit(name.c_str()) != 0) { | |
| 83 | 1 | if (stbi_us* data = stbi_load_16(name.c_str(), &width, &height, &channels, 4)) { | |
| 84 | 1 | return Decode(data, width, height, 1.0f / 65535.0f, encoding); | |
| 85 | } | ||
| 86 | 9 | } else if (stbi_uc* data = stbi_load(name.c_str(), &width, &height, &channels, 4)) { | |
| 87 | 6 | return Decode(data, width, height, 1.0f / 255.0f, encoding); | |
| 88 | } | ||
| 89 | 12 | throw std::runtime_error("Cannot read " + name + ": " + stbi_failure_reason()); | |
| 90 | 7 | } | |
| 91 | |||
| 92 | } // namespace | ||
| 93 | |||
| 94 | 3 | glm::vec4 Image::Mean() const { | |
| 95 | 3 | glm::dvec4 sum(0.0); | |
| 96 | 7 | for (const glm::vec4& p : pixels) { | |
| 97 | 4 | sum += glm::dvec4(p); | |
| 98 | } | ||
| 99 | 3 | return pixels.empty() ? glm::vec4(0.0f) : glm::vec4(sum / static_cast<double>(pixels.size())); | |
| 100 | } | ||
| 101 | |||
| 102 | 3 | void WriteExr(const std::filesystem::path& path, const Tensor& image) { | |
| 103 | 3 | if ((image.Rank() != 2 && image.Rank() != 3) || image.Location().kind != DeviceKind::Cpu) { | |
| 104 | 1 | throw std::invalid_argument("EXR images are (rows, cols[, channels]) on the CPU"); | |
| 105 | } | ||
| 106 | 2 | const auto rows = static_cast<int>(image.Shape()[0]); | |
| 107 | 2 | const auto cols = static_cast<int>(image.Shape()[1]); | |
| 108 | 2 | const int channels = image.Rank() == 3 ? static_cast<int>(image.Shape()[2]) : 1; | |
| 109 | 2 | const DType type = image.Type(); | |
| 110 | 2 | if (type != DType::Float32 && type != DType::UInt32 && type != DType::UInt16 && | |
| 111 | type != DType::UInt8) { | ||
| 112 | ✗ | throw std::invalid_argument("EXR images are float32 or unsigned integers"); | |
| 113 | } | ||
| 114 | 2 | const Tensor contiguous = image.IsContiguous() ? image : image.Contiguous(); | |
| 115 | 2 | const std::size_t count = static_cast<std::size_t>(rows) * static_cast<std::size_t>(cols) * | |
| 116 | 2 | static_cast<std::size_t>(channels); | |
| 117 | // Floats and uint32 as they are; narrower integers as floats. | ||
| 118 | 2 | const bool uints = type == DType::UInt32; | |
| 119 | 2 | std::vector<float> floats; | |
| 120 | 2 | if (type == DType::UInt16 || type == DType::UInt8) { | |
| 121 | ✗ | floats.resize(count); | |
| 122 | ✗ | for (std::size_t i = 0; i < count; ++i) { | |
| 123 | ✗ | floats[i] = | |
| 124 | ✗ | static_cast<float>(type == DType::UInt16 ? contiguous.As<std::uint16_t>()[i] | |
| 125 | ✗ | : contiguous.As<std::uint8_t>()[i]); | |
| 126 | } | ||
| 127 | } | ||
| 128 | 2 | char* base = floats.empty() ? static_cast<char*>(const_cast<void*>(contiguous.Data())) | |
| 129 | ✗ | : reinterpret_cast<char*>(floats.data()); | |
| 130 | 2 | static constexpr std::array<const char*, 4> kRgba{"R", "G", "B", "A"}; | |
| 131 | 12 | const auto name = [&](int c) -> std::string { | |
| 132 | 10 | if (channels == 1) { | |
| 133 | 2 | return "Y"; | |
| 134 | } | ||
| 135 | 8 | return channels <= 4 ? kRgba[static_cast<std::size_t>(c)] : std::to_string(c); | |
| 136 | 2 | }; | |
| 137 | 2 | try { | |
| 138 | 2 | Imf::Header header(cols, rows); | |
| 139 | 2 | header.compression() = Imf::ZIP_COMPRESSION; | |
| 140 | 2 | const Imf::PixelType pixel = uints ? Imf::UINT : Imf::FLOAT; | |
| 141 | 2 | Imf::FrameBuffer frame; | |
| 142 | 2 | const std::size_t texel = static_cast<std::size_t>(channels) * sizeof(float); | |
| 143 | 7 | for (int c = 0; c < channels; ++c) { | |
| 144 | 5 | header.channels().insert(name(c), Imf::Channel(pixel)); | |
| 145 | 5 | frame.insert(name(c), | |
| 146 | 10 | Imf::Slice(pixel, base + static_cast<std::size_t>(c) * sizeof(float), | |
| 147 | texel, texel * static_cast<std::size_t>(cols))); | ||
| 148 | } | ||
| 149 | 2 | Imf::OutputFile file(path.string().c_str(), header); | |
| 150 | 2 | file.setFrameBuffer(frame); | |
| 151 | 2 | file.writePixels(rows); | |
| 152 | 2 | } catch (const std::exception& error) { | |
| 153 | ✗ | throw std::runtime_error("Cannot write " + path.string() + ": " + error.what()); | |
| 154 | ✗ | } | |
| 155 | 2 | } | |
| 156 | |||
| 157 | 20 | Image LoadImageFile(const std::filesystem::path& path, ImageEncoding encoding) { | |
| 158 | 20 | ERNIE_ZONE("Decode image"); | |
| 159 | 60 | const std::string extension = ToLower(path.extension().string()); | |
| 160 | 20 | if (extension == ".exr") { | |
| 161 | 9 | return LoadExr(path); | |
| 162 | } | ||
| 163 | 11 | if (IsPngOrJpeg(path)) { | |
| 164 | 10 | return LoadStb(path, encoding); | |
| 165 | } | ||
| 166 | 2 | throw std::runtime_error("Unsupported image format: " + path.string()); | |
| 167 | 20 | } | |
| 168 | |||
| 169 | 61 | bool IsPngOrJpeg(const std::filesystem::path& path) { | |
| 170 | 183 | const std::string extension = ToLower(path.extension().string()); | |
| 171 | 61 | return extension == ".png" || extension == ".jpg" || extension == ".jpeg"; | |
| 172 | 61 | } | |
| 173 | |||
| 174 | 4633 | float SrgbToLinear(float value) { | |
| 175 | 4633 | return value <= 0.04045f ? value / 12.92f : std::pow((value + 0.055f) / 1.055f, 2.4f); | |
| 176 | } | ||
| 177 | |||
| 178 | 134868 | float LinearToSrgb(float value) { | |
| 179 | 134868 | return value <= 0.0031308f ? value * 12.92f : 1.055f * std::pow(value, 1.0f / 2.4f) - 0.055f; | |
| 180 | } | ||
| 181 | |||
| 182 | 346 | float EncodeSlopeDeviation(float sigma) { | |
| 183 | 346 | const float s = std::max(sigma, 0.0f); | |
| 184 | 346 | return s / (1.0f + s); | |
| 185 | } | ||
| 186 | |||
| 187 | 1383 | float DecodeSlopeDeviation(float encoded) { | |
| 188 | 1383 | const float e = std::clamp(encoded, 0.0f, 0.999f); | |
| 189 | 1383 | return e / (1.0f - e); | |
| 190 | } | ||
| 191 | |||
| 192 | namespace { | ||
| 193 | |||
| 194 | /// The slope (dz/dx, dz/dy) of a normal stored as rgb in [0, 1]. | ||
| 195 | 1380 | glm::vec2 NormalSlope(const glm::vec4& texel) { | |
| 196 | 1380 | const glm::vec3 n = 2.0f * glm::vec3(texel) - 1.0f; | |
| 197 | 1380 | return -glm::vec2(n) / std::max(n.z, kMinNormalZ); | |
| 198 | } | ||
| 199 | |||
| 200 | /// The 2 x 2 slope-space filter of `NormalSlopes`: the mean slope's | ||
| 201 | /// normal, and the slopes' deviation from it, within and between the | ||
| 202 | /// four texels, per axis. | ||
| 203 | 345 | glm::vec4 FilterSlopes(const std::array<glm::vec4, 4>& texels) { | |
| 204 | 345 | glm::vec2 mean(0.0f); | |
| 205 | 345 | float square = 0.0f; // Mean of sx^2 + sy^2 over everything the | |
| 206 | // texels cover. | ||
| 207 | 1725 | for (const glm::vec4& texel : texels) { | |
| 208 | 1380 | const glm::vec2 slope = NormalSlope(texel); | |
| 209 | 1380 | const float sigma = DecodeSlopeDeviation(texel.a); | |
| 210 | 1380 | mean += 0.25f * slope; | |
| 211 | 1380 | square += 0.25f * (glm::dot(slope, slope) + 2.0f * sigma * sigma); | |
| 212 | } | ||
| 213 | 345 | const float variance = std::max(0.5f * (square - glm::dot(mean, mean)), 0.0f); | |
| 214 | 345 | const glm::vec3 n = glm::normalize(glm::vec3(-mean, 1.0f)); | |
| 215 | 345 | return glm::vec4(0.5f * n + 0.5f, EncodeSlopeDeviation(std::sqrt(variance))); | |
| 216 | } | ||
| 217 | |||
| 218 | } // namespace | ||
| 219 | |||
| 220 | 11 | std::vector<Image> MipChain(const Image& image, MipFilter filter) { | |
| 221 | 11 | ERNIE_ZONE("Mip chain"); | |
| 222 | 22 | std::vector<Image> levels{image}; | |
| 223 | 11 | if (filter == MipFilter::NormalSlopes) { | |
| 224 | 1039 | for (glm::vec4& texel : levels.back().pixels) { | |
| 225 | 1036 | texel.a = 0.0f; | |
| 226 | } | ||
| 227 | } | ||
| 228 | 37 | while (levels.back().width > 1 || levels.back().height > 1) { | |
| 229 | 26 | const Image& above = levels.back(); | |
| 230 | 53 | Image level(std::max(1u, above.width / 2), std::max(1u, above.height / 2)); | |
| 231 | 105 | for (std::uint32_t y = 0; y < level.height; ++y) { | |
| 232 | 79 | const std::uint32_t y0 = std::min(2 * y, above.height - 1); | |
| 233 | 79 | const std::uint32_t y1 = std::min(2 * y + 1, above.height - 1); | |
| 234 | 732 | for (std::uint32_t x = 0; x < level.width; ++x) { | |
| 235 | 653 | const std::uint32_t x0 = std::min(2 * x, above.width - 1); | |
| 236 | 653 | const std::uint32_t x1 = std::min(2 * x + 1, above.width - 1); | |
| 237 | 653 | if (filter == MipFilter::NormalSlopes) { | |
| 238 | 345 | level.At(x, y) = FilterSlopes( | |
| 239 | 345 | {above.At(x0, y0), above.At(x1, y0), above.At(x0, y1), above.At(x1, y1)}); | |
| 240 | 345 | continue; | |
| 241 | } | ||
| 242 | 308 | glm::vec4 texel = 0.25f * (above.At(x0, y0) + above.At(x1, y0) + above.At(x0, y1) + | |
| 243 | 616 | above.At(x1, y1)); | |
| 244 | 308 | if (filter == MipFilter::UnitVector) { | |
| 245 | 2 | glm::vec3 n = 2.0f * glm::vec3(texel) - 1.0f; | |
| 246 | 2 | const float length = glm::length(n); | |
| 247 | 2 | n = length > 0.0f ? n / length : glm::vec3(0.0f, 0.0f, 1.0f); | |
| 248 | 2 | texel = glm::vec4(0.5f * n + 0.5f, texel.a); | |
| 249 | } | ||
| 250 | 308 | level.At(x, y) = texel; | |
| 251 | } | ||
| 252 | } | ||
| 253 | 26 | levels.push_back(std::move(level)); | |
| 254 | 26 | } | |
| 255 | 11 | return levels; | |
| 256 | 22 | } | |
| 257 | |||
| 258 | 39 | std::vector<std::uint8_t> ToRgba8(const Image& image, ImageEncoding encoding) { | |
| 259 | 39 | ERNIE_ZONE("Quantise to RGBA8"); | |
| 260 | 39 | std::vector<std::uint8_t> bytes(image.pixels.size() * 4); | |
| 261 | 13683 | auto quantise = [](float value) { | |
| 262 | 13644 | return static_cast<std::uint8_t>(std::lround(std::clamp(value, 0.0f, 1.0f) * 255.0f)); | |
| 263 | }; | ||
| 264 | 3450 | for (std::size_t i = 0; i < image.pixels.size(); ++i) { | |
| 265 | 3411 | const glm::vec4& p = image.pixels[i]; | |
| 266 | 13644 | for (int c = 0; c < 3; ++c) { | |
| 267 | 10233 | const float value = encoding == ImageEncoding::Srgb ? LinearToSrgb(p[c]) : p[c]; | |
| 268 | 10233 | bytes[4 * i + static_cast<std::size_t>(c)] = quantise(value); | |
| 269 | } | ||
| 270 | 3411 | bytes[4 * i + 3] = quantise(p.a); | |
| 271 | } | ||
| 272 | 39 | return bytes; | |
| 273 | 39 | } | |
| 274 | |||
| 275 | ✗ | Image ImageFromTensor(const Tensor& tensor, ImageEncoding encoding) { | |
| 276 | ✗ | const bool grey = tensor.Rank() == 2; | |
| 277 | ✗ | const std::int64_t channels = grey ? 1 : (tensor.Rank() == 3 ? tensor.Shape()[2] : 0); | |
| 278 | ✗ | const bool typed = tensor.Type() == DType::Float32 || tensor.Type() == DType::UInt8; | |
| 279 | ✗ | if (!typed || (channels != 1 && channels != 3 && channels != 4) || | |
| 280 | ✗ | tensor.Location().kind != DeviceKind::Cpu || tensor.Shape()[0] < 1 || | |
| 281 | ✗ | tensor.Shape()[1] < 1) { | |
| 282 | ✗ | throw std::invalid_argument("expected a float32 or uint8 (rows, cols[, 1, 3 or 4]) " | |
| 283 | ✗ | "CPU tensor"); | |
| 284 | } | ||
| 285 | ✗ | const Tensor packed = tensor.Contiguous(); | |
| 286 | ✗ | Image image(static_cast<std::uint32_t>(tensor.Shape()[1]), | |
| 287 | ✗ | static_cast<std::uint32_t>(tensor.Shape()[0])); | |
| 288 | ✗ | const bool bytes = tensor.Type() == DType::UInt8; | |
| 289 | ✗ | const auto* floats = bytes ? nullptr : packed.As<float>(); | |
| 290 | ✗ | const auto* uints = bytes ? packed.As<std::uint8_t>() : nullptr; | |
| 291 | ✗ | const auto value = [&](std::size_t index, bool colour) { | |
| 292 | ✗ | if (!bytes) { | |
| 293 | ✗ | return floats[index]; | |
| 294 | } | ||
| 295 | ✗ | const float stored = static_cast<float>(uints[index]) / 255.0f; | |
| 296 | ✗ | return colour && encoding == ImageEncoding::Srgb ? SrgbToLinear(stored) : stored; | |
| 297 | ✗ | }; | |
| 298 | ✗ | const auto stride = static_cast<std::size_t>(channels); | |
| 299 | ✗ | for (std::size_t i = 0; i < image.pixels.size(); ++i) { | |
| 300 | ✗ | glm::vec4& pixel = image.pixels[i]; | |
| 301 | ✗ | const std::size_t base = i * stride; | |
| 302 | ✗ | if (channels == 1) { | |
| 303 | ✗ | pixel = glm::vec4(glm::vec3(value(base, true)), 1.0f); | |
| 304 | } else { | ||
| 305 | ✗ | pixel = glm::vec4(value(base, true), value(base + 1, true), value(base + 2, true), | |
| 306 | ✗ | channels == 4 ? value(base + 3, false) : 1.0f); | |
| 307 | } | ||
| 308 | } | ||
| 309 | ✗ | return image; | |
| 310 | ✗ | } | |
| 311 | |||
| 312 | } // namespace ernie | ||
| 313 |