io/image.cpp
| Line | Branch | Exec | Source |
|---|---|---|---|
| 1 | // SPDX-License-Identifier: BSD-3-Clause | ||
| 2 | #include "chrono_batch_render/io/image.h" | ||
| 3 | |||
| 4 | #include <ImfChannelList.h> | ||
| 5 | #include <ImfFrameBuffer.h> | ||
| 6 | #include <ImfHeader.h> | ||
| 7 | #include <ImfIO.h> | ||
| 8 | #include <ImfOutputFile.h> | ||
| 9 | #include <ImfRgbaFile.h> | ||
| 10 | #include <stb_image.h> | ||
| 11 | #include <stb_image_write.h> | ||
| 12 | |||
| 13 | #include <algorithm> | ||
| 14 | #include <array> | ||
| 15 | #include <cmath> | ||
| 16 | #include <cstring> | ||
| 17 | #include <fstream> | ||
| 18 | #include <iterator> | ||
| 19 | #include <limits> | ||
| 20 | #include <memory> | ||
| 21 | #include <span> | ||
| 22 | #include <string> | ||
| 23 | #include <string_view> | ||
| 24 | #include <tuple> | ||
| 25 | #include <utility> | ||
| 26 | #include <vector> | ||
| 27 | |||
| 28 | #include "chrono_batch_render/device/buffer.h" | ||
| 29 | #include "chrono_batch_render/device/commands.h" | ||
| 30 | #include "chrono_batch_render/device/device.h" | ||
| 31 | #include "chrono_batch_render/device/program.h" | ||
| 32 | #include "chrono_batch_render/device/queue.h" | ||
| 33 | #include "chrono_batch_render/device/slang/dispatch.h" | ||
| 34 | #include "chrono_batch_render/io/kernels.h" | ||
| 35 | #include "chrono_batch_render/io/kernels/colorize.h" | ||
| 36 | #include "chrono_batch_render/io/kernels/encoding.h" | ||
| 37 | #include "chrono_batch_render/util/error.h" | ||
| 38 | #include "chrono_batch_render/util/profiler.h" | ||
| 39 | #include "chrono_batch_render/util/span.h" | ||
| 40 | #include "chrono_batch_render/util/text.h" | ||
| 41 | |||
| 42 | namespace chrono::batch_render { | ||
| 43 | |||
| 44 | namespace { | ||
| 45 | |||
| 46 | 1 | class MemoryStream : public Imf::IStream { | |
| 47 | public: | ||
| 48 | 8 | MemoryStream(std::span<const std::byte> bytes, const std::string& name) | |
| 49 | 16 | : Imf::IStream(name.c_str()), m_bytes(bytes) {} | |
| 50 | |||
| 51 | ✗ | bool read(char c[], int n) override { | |
| 52 | ✗ | if (m_at + uint64_t(n) > m_bytes.size()) { | |
| 53 | ✗ | throw Error("{} ends early", fileName()); | |
| 54 | } | ||
| 55 | ✗ | std::memcpy(c, m_bytes.data() + m_at, size_t(n)); | |
| 56 | ✗ | m_at += uint64_t(n); | |
| 57 | ✗ | return m_at < m_bytes.size(); | |
| 58 | } | ||
| 59 | |||
| 60 | ✗ | uint64_t tellg() override { | |
| 61 | ✗ | return m_at; | |
| 62 | } | ||
| 63 | |||
| 64 | ✗ | void seekg(uint64_t at) override { | |
| 65 | ✗ | m_at = at; | |
| 66 | ✗ | } | |
| 67 | |||
| 68 | 8 | int64_t size() override { | |
| 69 | 8 | return int64_t(m_bytes.size()); | |
| 70 | } | ||
| 71 | |||
| 72 | 8 | bool isStatelessRead() const override { | |
| 73 | 8 | return true; | |
| 74 | } | ||
| 75 | |||
| 76 | 36 | int64_t read(void* buffer, uint64_t count, uint64_t offset) override { | |
| 77 | 36 | if (offset >= m_bytes.size()) { | |
| 78 | return 0; | ||
| 79 | } | ||
| 80 | 36 | const uint64_t available = std::min<uint64_t>(count, m_bytes.size() - offset); | |
| 81 | 36 | std::memcpy(buffer, m_bytes.data() + offset, size_t(available)); | |
| 82 | 36 | return int64_t(available); | |
| 83 | } | ||
| 84 | |||
| 85 | private: | ||
| 86 | std::span<const std::byte> m_bytes; | ||
| 87 | uint64_t m_at = 0; | ||
| 88 | }; | ||
| 89 | |||
| 90 | 8 | Image LoadExr(std::span<const std::byte> bytes, const std::string& name) { | |
| 91 | 8 | try { | |
| 92 | 8 | MemoryStream stream(bytes, name); | |
| 93 | 8 | Imf::RgbaInputFile file(stream); | |
| 94 | 7 | const Imath::Box2i window = file.dataWindow(); | |
| 95 | 7 | const int width = window.max.x - window.min.x + 1; | |
| 96 | 7 | const int height = window.max.y - window.min.y + 1; | |
| 97 | 7 | std::vector<Imf::Rgba> texels(size_t(width) * size_t(height)); | |
| 98 | 7 | file.setFrameBuffer(texels.data() - window.min.x - std::ptrdiff_t(window.min.y) * width, 1, | |
| 99 | size_t(width)); | ||
| 100 | 7 | file.readPixels(window.min.y, window.max.y); | |
| 101 | |||
| 102 | 7 | Image image{uint32_t(width), uint32_t(height)}; | |
| 103 | 85 | for (size_t i = 0; i < image.pixels.size(); ++i) { | |
| 104 | 78 | const Imf::Rgba& t = texels[i]; | |
| 105 | 78 | image.pixels[i] = {float(t.r), float(t.g), float(t.b), float(t.a)}; | |
| 106 | } | ||
| 107 | 7 | return image; | |
| 108 | 8 | } catch (const std::exception& error) { | |
| 109 | 1 | throw Error("Cannot read {}: {}", name, error.what()); | |
| 110 | 1 | } | |
| 111 | } | ||
| 112 | |||
| 113 | template <class T> | ||
| 114 | 6 | Image Decode(T* data, int width, int height, float scale, ImageEncoding encoding) { | |
| 115 | 6 | Image image{uint32_t(width), uint32_t(height)}; | |
| 116 | 58 | for (size_t i = 0; i < image.pixels.size(); ++i) { | |
| 117 | 52 | Rgba value{float(data[4 * i]) * scale, float(data[4 * i + 1]) * scale, | |
| 118 | 52 | float(data[4 * i + 2]) * scale, float(data[4 * i + 3]) * scale}; | |
| 119 | 52 | if (encoding == ImageEncoding::Srgb) { | |
| 120 | 12 | value = {DecodeSrgb(value.r), DecodeSrgb(value.g), DecodeSrgb(value.b), value.a}; | |
| 121 | } | ||
| 122 | 52 | image.pixels[i] = value; | |
| 123 | } | ||
| 124 | 6 | stbi_image_free(data); | |
| 125 | 6 | return image; | |
| 126 | ✗ | } | |
| 127 | |||
| 128 | 7 | Image LoadStb(std::span<const std::byte> bytes, const std::string& name, ImageEncoding encoding) { | |
| 129 | 7 | const auto* data = reinterpret_cast<const stbi_uc*>(bytes.data()); | |
| 130 | 7 | const int size = int(bytes.size()); | |
| 131 | 7 | int width = 0; | |
| 132 | 7 | int height = 0; | |
| 133 | 7 | int channels = 0; | |
| 134 | 7 | if (stbi_is_16_bit_from_memory(data, size) != 0) { | |
| 135 | ✗ | if (stbi_us* pixels = stbi_load_16_from_memory(data, size, &width, &height, &channels, 4)) { | |
| 136 | ✗ | return Decode(pixels, width, height, 1.0f / 65535.0f, encoding); | |
| 137 | } | ||
| 138 | 7 | } else if (stbi_uc* pixels = stbi_load_from_memory(data, size, &width, &height, &channels, 4)) { | |
| 139 | 6 | return Decode(pixels, width, height, 1.0f / 255.0f, encoding); | |
| 140 | } | ||
| 141 | 1 | throw Error("Cannot read {}: {}", name, stbi_failure_reason()); | |
| 142 | } | ||
| 143 | |||
| 144 | 1 | Image ToImage(const Pixels& pixels, ImageEncoding encoding) { | |
| 145 | 1 | const bool bytes = pixels.type == device::DType::UInt8; | |
| 146 | 1 | if ((!bytes && pixels.type != device::DType::Float32) || | |
| 147 | 1 | (pixels.channels != 1 && pixels.channels != 3 && pixels.channels != 4)) { | |
| 148 | ✗ | throw Error("Images are made of float32 or uint8 (rows, cols[, 1, 3 or 4]) pixels"); | |
| 149 | } | ||
| 150 | |||
| 151 | 1 | Image image{uint32_t(pixels.cols), uint32_t(pixels.rows)}; | |
| 152 | 25 | const auto value = [&](size_t index, bool color) { | |
| 153 | 24 | if (!bytes) { | |
| 154 | ✗ | float stored; | |
| 155 | ✗ | std::memcpy(&stored, pixels.bytes.data() + index * sizeof(float), sizeof(float)); | |
| 156 | ✗ | return stored; | |
| 157 | } | ||
| 158 | 24 | const float stored = float(uint8_t(pixels.bytes[index])) / 255.0f; | |
| 159 | 24 | return color && encoding == ImageEncoding::Srgb ? DecodeSrgb(stored) : stored; | |
| 160 | 1 | }; | |
| 161 | |||
| 162 | 7 | for (size_t i = 0; i < image.pixels.size(); ++i) { | |
| 163 | 6 | const size_t base = i * size_t(pixels.channels); | |
| 164 | 6 | if (pixels.channels == 1) { | |
| 165 | ✗ | const float gray = value(base, true); | |
| 166 | ✗ | image.pixels[i] = {gray, gray, gray, 1}; | |
| 167 | } else { | ||
| 168 | 12 | image.pixels[i] = {value(base, true), value(base + 1, true), value(base + 2, true), | |
| 169 | 6 | pixels.channels == 4 ? value(base + 3, false) : 1}; | |
| 170 | } | ||
| 171 | } | ||
| 172 | 1 | return image; | |
| 173 | } | ||
| 174 | |||
| 175 | 23 | std::vector<uint8_t> Encoded(const Pixels& pixels, const char* format) { | |
| 176 | 23 | const int64_t channels = pixels.channels; | |
| 177 | 23 | if (pixels.type == device::DType::UInt8 && channels >= 1 && channels <= 4) { | |
| 178 | 21 | std::vector<uint8_t> bytes(pixels.bytes.size()); | |
| 179 | 21 | std::memcpy(bytes.data(), pixels.bytes.data(), bytes.size()); | |
| 180 | 21 | return bytes; | |
| 181 | 21 | } | |
| 182 | 2 | if (pixels.type == device::DType::UInt32) { | |
| 183 | 1 | throw Error("{} images can't hold uint32 ids, colorize them with ColorizeIds or write .exr", | |
| 184 | 1 | format); | |
| 185 | } | ||
| 186 | 1 | if (pixels.type == device::DType::Float32 && channels == 1) { | |
| 187 | 1 | throw Error("{} images can't hold float32 single channel data, colorize it with " | |
| 188 | "ColorizeScalars or write .exr", | ||
| 189 | 1 | format); | |
| 190 | } | ||
| 191 | ✗ | if (pixels.type != device::DType::Float32 || (channels != 3 && channels != 4)) { | |
| 192 | ✗ | throw Error("{} images are uint8 (rows, cols[, 1 to 4]) or float32 linear color (rows, " | |
| 193 | "cols, 3 or 4) pixels", | ||
| 194 | ✗ | format); | |
| 195 | } | ||
| 196 | |||
| 197 | ✗ | const std::vector<uint8_t> rgba = | |
| 198 | ✗ | ToRgba8(ToImage(pixels, ImageEncoding::Linear), ImageEncoding::Srgb); | |
| 199 | ✗ | const size_t count = size_t(pixels.rows) * size_t(pixels.cols); | |
| 200 | ✗ | std::vector<uint8_t> bytes(count * size_t(channels)); | |
| 201 | ✗ | for (size_t i = 0; i < count; ++i) { | |
| 202 | ✗ | for (size_t c = 0; c < size_t(channels); ++c) { | |
| 203 | ✗ | bytes[i * channels + c] = rgba[4 * i + c]; | |
| 204 | } | ||
| 205 | } | ||
| 206 | ✗ | return bytes; | |
| 207 | ✗ | } | |
| 208 | |||
| 209 | 46 | void WritePngBytes(const std::filesystem::path& path, const uint8_t* bytes, int rows, int cols, | |
| 210 | int channels) { | ||
| 211 | 46 | if (stbi_write_png(path.string().c_str(), cols, rows, channels, bytes, cols * channels) == 0) { | |
| 212 | ✗ | throw Error("Cannot write {}", path.string()); | |
| 213 | } | ||
| 214 | 46 | } | |
| 215 | |||
| 216 | 1 | void WriteJpegBytes(const std::filesystem::path& path, const uint8_t* bytes, int rows, int cols, | |
| 217 | int channels, int quality) { | ||
| 218 | 1 | if (quality < 1 || quality > 100) { | |
| 219 | ✗ | throw Error("JPEG quality is from 1 to 100, not {}", quality); | |
| 220 | } | ||
| 221 | 1 | if (stbi_write_jpg(path.string().c_str(), cols, rows, channels, bytes, quality) == 0) { | |
| 222 | ✗ | throw Error("Cannot write {}", path.string()); | |
| 223 | } | ||
| 224 | 1 | } | |
| 225 | |||
| 226 | 10 | void WriteExrChannels(const std::filesystem::path& path, char* base, int rows, int cols, | |
| 227 | int channels, Imf::PixelType pixel) { | ||
| 228 | 10 | static constexpr std::array<const char*, 4> kRgba{"R", "G", "B", "A"}; | |
| 229 | 72 | const auto name = [&](int c) -> std::string { | |
| 230 | 62 | if (channels == 1) { | |
| 231 | 6 | return "Y"; | |
| 232 | } | ||
| 233 | 56 | return channels <= 4 ? kRgba[size_t(c)] : std::to_string(c); | |
| 234 | 10 | }; | |
| 235 | |||
| 236 | 10 | try { | |
| 237 | 10 | Imf::Header header(cols, rows); | |
| 238 | 10 | header.compression() = Imf::ZIP_COMPRESSION; | |
| 239 | 10 | Imf::FrameBuffer frame; | |
| 240 | 10 | const size_t texel = size_t(channels) * sizeof(float); | |
| 241 | 41 | for (int c = 0; c < channels; ++c) { | |
| 242 | 31 | header.channels().insert(name(c), Imf::Channel(pixel)); | |
| 243 | 62 | frame.insert(name(c), Imf::Slice(pixel, base + size_t(c) * sizeof(float), texel, | |
| 244 | 31 | texel * size_t(cols))); | |
| 245 | } | ||
| 246 | |||
| 247 | 10 | Imf::OutputFile file(path.string().c_str(), header); | |
| 248 | 10 | file.setFrameBuffer(frame); | |
| 249 | 10 | file.writePixels(rows); | |
| 250 | 10 | } catch (const std::exception& error) { | |
| 251 | ✗ | throw Error("Cannot write {}: {}", path.string(), error.what()); | |
| 252 | ✗ | } | |
| 253 | 10 | } | |
| 254 | |||
| 255 | } | ||
| 256 | |||
| 257 | 1 | Rgba Image::Mean() const { | |
| 258 | 1 | if (pixels.empty()) { | |
| 259 | ✗ | return {}; | |
| 260 | } | ||
| 261 | 1 | double sum[4] = {}; | |
| 262 | 13 | for (const Rgba& p : pixels) { | |
| 263 | 12 | sum[0] += p.r; | |
| 264 | 12 | sum[1] += p.g; | |
| 265 | 12 | sum[2] += p.b; | |
| 266 | 12 | sum[3] += p.a; | |
| 267 | } | ||
| 268 | 1 | const double count = double(pixels.size()); | |
| 269 | 1 | return {float(sum[0] / count), float(sum[1] / count), float(sum[2] / count), | |
| 270 | 1 | float(sum[3] / count)}; | |
| 271 | } | ||
| 272 | |||
| 273 | namespace { | ||
| 274 | |||
| 275 | 31 | std::vector<char> ReadBytes(const std::filesystem::path& path) { | |
| 276 | 31 | std::ifstream file(path, std::ios::binary); | |
| 277 | 31 | if (!file) { | |
| 278 | 16 | throw Error("Cannot read {}: can't open it", path.string()); | |
| 279 | } | ||
| 280 | 46 | return {std::istreambuf_iterator<char>(file), std::istreambuf_iterator<char>()}; | |
| 281 | 31 | } | |
| 282 | |||
| 283 | } | ||
| 284 | |||
| 285 | 11 | Image LoadImageFile(const std::filesystem::path& path, ImageEncoding encoding) { | |
| 286 | 11 | const std::vector<char> text = ReadBytes(path); | |
| 287 | 18 | return LoadImageBytes(std::as_bytes(std::span(text)), path.string(), encoding); | |
| 288 | 9 | } | |
| 289 | |||
| 290 | 20 | Rgba8Image LoadRgba8File(const std::filesystem::path& path) { | |
| 291 | 20 | const std::vector<char> text = ReadBytes(path); | |
| 292 | 28 | return LoadRgba8Bytes(std::as_bytes(std::span(text)), path.string()); | |
| 293 | 14 | } | |
| 294 | |||
| 295 | 23 | Rgba8Image LoadRgba8Bytes(Span<const std::byte> bytes, std::string_view name) { | |
| 296 | 23 | CHRONO_BATCH_RENDER_ZONE("IO/LoadImage"); | |
| 297 | 23 | int width = 0; | |
| 298 | 23 | int height = 0; | |
| 299 | 23 | int channels = 0; | |
| 300 | 23 | stbi_uc* pixels = stbi_load_from_memory(reinterpret_cast<const stbi_uc*>(bytes.data()), | |
| 301 | 23 | int(bytes.size()), &width, &height, &channels, 4); | |
| 302 | 23 | if (pixels == nullptr) { | |
| 303 | 1 | throw Error("Cannot read {}: {}", name, stbi_failure_reason()); | |
| 304 | } | ||
| 305 | 22 | Rgba8Image image{uint32_t(width), uint32_t(height), {}}; | |
| 306 | 22 | image.texels.resize(size_t(width) * size_t(height)); | |
| 307 | 22 | std::memcpy(image.texels.data(), pixels, image.texels.size() * sizeof(uint32_t)); | |
| 308 | 22 | stbi_image_free(pixels); | |
| 309 | 22 | return image; | |
| 310 | 23 | } | |
| 311 | |||
| 312 | 16 | Image LoadImageBytes(Span<const std::byte> bytes, std::string_view name, ImageEncoding encoding) { | |
| 313 | 16 | CHRONO_BATCH_RENDER_ZONE("IO/LoadImage"); | |
| 314 | 16 | const std::string named(name); | |
| 315 | 64 | if (Lowercase(std::filesystem::path(named).extension().string()) == ".exr") { | |
| 316 | 8 | return LoadExr(bytes, named); | |
| 317 | } | ||
| 318 | 16 | if (IsPngOrJpeg(named)) { | |
| 319 | 9 | return LoadStb(bytes, named, encoding); | |
| 320 | } | ||
| 321 | 1 | throw Error("Unsupported image format: {}", named); | |
| 322 | 16 | } | |
| 323 | |||
| 324 | 107 | bool IsPngOrJpeg(const std::filesystem::path& path) { | |
| 325 | 321 | const std::string extension = Lowercase(path.extension().string()); | |
| 326 | 107 | return extension == ".png" || extension == ".jpg" || extension == ".jpeg"; | |
| 327 | 107 | } | |
| 328 | |||
| 329 | 22 | Pixels::Pixels(const device::Tensor& tensor) : type(tensor.Type()) { | |
| 330 | 22 | const Span<const int64_t> shape = tensor.Shape(); | |
| 331 | 22 | if ((shape.size() != 2 && shape.size() != 3) || shape[0] < 1 || shape[1] < 1) { | |
| 332 | ✗ | throw Error("Pixels come from (rows, cols[, channels]) tensors, not {} dimensions", | |
| 333 | ✗ | shape.size()); | |
| 334 | } | ||
| 335 | 22 | rows = shape[0]; | |
| 336 | 22 | cols = shape[1]; | |
| 337 | 22 | channels = shape.size() == 3 ? shape[2] : 1; | |
| 338 | 22 | bytes = tensor.ToHostBytes(); | |
| 339 | 22 | } | |
| 340 | |||
| 341 | 30 | Pixels::Pixels(Span<const std::byte> data, device::DType elementType, int64_t rowCount, | |
| 342 | 30 | int64_t colCount, int64_t channelCount) | |
| 343 | 30 | : rows(rowCount), cols(colCount), channels(channelCount), type(elementType), | |
| 344 | 30 | bytes(data.begin(), data.end()) { | |
| 345 | 30 | if (rows < 1 || cols < 1 || channels < 1) { | |
| 346 | ✗ | throw Error("Pixels are at least 1 by 1 by 1, not {} by {} by {}", rows, cols, channels); | |
| 347 | } | ||
| 348 | 30 | const uint64_t expected = | |
| 349 | 30 | uint64_t(rows) * uint64_t(cols) * uint64_t(channels) * device::ElementBytes(type); | |
| 350 | 30 | if (data.size() != expected) { | |
| 351 | 1 | throw Error("{} by {} by {} pixels are {} bytes, not {}", rows, cols, channels, expected, | |
| 352 | 1 | data.size()); | |
| 353 | } | ||
| 354 | 30 | } | |
| 355 | |||
| 356 | 1 | Image ImageFromTensor(const device::Tensor& tensor, ImageEncoding encoding) { | |
| 357 | 2 | return ToImage(Pixels(tensor), encoding); | |
| 358 | } | ||
| 359 | |||
| 360 | 26 | std::vector<uint8_t> ToRgba8(const Image& image, ImageEncoding encoding) { | |
| 361 | 26 | std::vector<uint8_t> bytes(image.pixels.size() * 4); | |
| 362 | 16781030 | const auto quantize = [](float value) { | |
| 363 | 16781004 | return uint8_t(std::lround(std::clamp(value, 0.0f, 1.0f) * 255.0f)); | |
| 364 | }; | ||
| 365 | 12585779 | const auto color = [&](float value) { | |
| 366 | 12585753 | return quantize(encoding == ImageEncoding::Srgb ? EncodeSrgb(value) : value); | |
| 367 | 26 | }; | |
| 368 | |||
| 369 | 4195277 | for (size_t i = 0; i < image.pixels.size(); ++i) { | |
| 370 | 4195251 | const Rgba& p = image.pixels[i]; | |
| 371 | 4195251 | bytes[4 * i] = color(p.r); | |
| 372 | 4195251 | bytes[4 * i + 1] = color(p.g); | |
| 373 | 4195251 | bytes[4 * i + 2] = color(p.b); | |
| 374 | 4195251 | bytes[4 * i + 3] = quantize(p.a); | |
| 375 | } | ||
| 376 | 26 | return bytes; | |
| 377 | } | ||
| 378 | |||
| 379 | 26 | void WritePng(const std::filesystem::path& path, const Image& image) { | |
| 380 | 26 | CHRONO_BATCH_RENDER_ZONE("IO/WritePng"); | |
| 381 | 26 | const std::vector<uint8_t> bytes = ToRgba8(image, ImageEncoding::Srgb); | |
| 382 | 26 | WritePngBytes(path, bytes.data(), int(image.height), int(image.width), 4); | |
| 383 | 26 | } | |
| 384 | |||
| 385 | 22 | void WritePng(const std::filesystem::path& path, const Pixels& image) { | |
| 386 | 22 | CHRONO_BATCH_RENDER_ZONE("IO/WritePng"); | |
| 387 | 22 | const std::vector<uint8_t> bytes = Encoded(image, "PNG"); | |
| 388 | 20 | WritePngBytes(path, bytes.data(), int(image.rows), int(image.cols), int(image.channels)); | |
| 389 | 22 | } | |
| 390 | |||
| 391 | ✗ | void WriteJpeg(const std::filesystem::path& path, const Image& image, int quality) { | |
| 392 | ✗ | const std::vector<uint8_t> bytes = ToRgba8(image, ImageEncoding::Srgb); | |
| 393 | ✗ | WriteJpegBytes(path, bytes.data(), int(image.height), int(image.width), 4, quality); | |
| 394 | ✗ | } | |
| 395 | |||
| 396 | 1 | void WriteJpeg(const std::filesystem::path& path, const Pixels& image, int quality) { | |
| 397 | 1 | const std::vector<uint8_t> bytes = Encoded(image, "JPEG"); | |
| 398 | 1 | WriteJpegBytes(path, bytes.data(), int(image.rows), int(image.cols), int(image.channels), | |
| 399 | quality); | ||
| 400 | 1 | } | |
| 401 | |||
| 402 | 7 | void WriteExr(const std::filesystem::path& path, const Image& image) { | |
| 403 | 7 | CHRONO_BATCH_RENDER_ZONE("IO/WriteExr"); | |
| 404 | 7 | std::vector<Rgba> pixels = image.pixels; | |
| 405 | 7 | WriteExrChannels(path, reinterpret_cast<char*>(pixels.data()), int(image.height), | |
| 406 | 7 | int(image.width), 4, Imf::FLOAT); | |
| 407 | 7 | } | |
| 408 | |||
| 409 | 3 | void WriteExr(const std::filesystem::path& path, const Pixels& image) { | |
| 410 | 3 | CHRONO_BATCH_RENDER_ZONE("IO/WriteExr"); | |
| 411 | 3 | if (image.type != device::DType::Float32 && image.type != device::DType::UInt32 && | |
| 412 | image.type != device::DType::UInt8) { | ||
| 413 | ✗ | throw Error("EXR images are float32, uint32 or uint8 pixels"); | |
| 414 | } | ||
| 415 | |||
| 416 | 3 | std::vector<std::byte> bytes = image.bytes; | |
| 417 | 3 | std::vector<float> floats; | |
| 418 | 3 | if (image.type == device::DType::UInt8) { | |
| 419 | ✗ | floats.resize(bytes.size()); | |
| 420 | ✗ | for (size_t i = 0; i < floats.size(); ++i) { | |
| 421 | ✗ | floats[i] = float(uint8_t(bytes[i])); | |
| 422 | } | ||
| 423 | } | ||
| 424 | 3 | char* base = floats.empty() ? reinterpret_cast<char*>(bytes.data()) | |
| 425 | ✗ | : reinterpret_cast<char*>(floats.data()); | |
| 426 | 3 | WriteExrChannels(path, base, int(image.rows), int(image.cols), int(image.channels), | |
| 427 | 3 | image.type == device::DType::UInt32 ? Imf::UINT : Imf::FLOAT); | |
| 428 | 3 | } | |
| 429 | |||
| 430 | 24 | std::filesystem::path WriteImage(std::filesystem::path path, const Pixels& image) { | |
| 431 | 24 | if (!path.has_extension()) { | |
| 432 | 9 | path += image.type == device::DType::UInt8 ? ".png" : ".exr"; | |
| 433 | } | ||
| 434 | 72 | const std::string extension = Lowercase(path.extension().string()); | |
| 435 | 24 | if (extension == ".png") { | |
| 436 | 19 | WritePng(path, image); | |
| 437 | 5 | } else if (extension == ".jpg" || extension == ".jpeg") { | |
| 438 | 1 | WriteJpeg(path, image); | |
| 439 | 4 | } else if (extension == ".exr") { | |
| 440 | 3 | WriteExr(path, image); | |
| 441 | } else { | ||
| 442 | 2 | throw Error("Cannot write {}: images are .png, .jpg, .jpeg or .exr", path.string()); | |
| 443 | } | ||
| 444 | 23 | return path; | |
| 445 | 23 | } | |
| 446 | |||
| 447 | namespace { | ||
| 448 | |||
| 449 | ✗ | device::Tensor ToTensor(const device::Tensor& like, const Pixels& pixels) { | |
| 450 | ✗ | device::Device& device = like.GetDevice(); | |
| 451 | ✗ | const device::BufferUsage usage = | |
| 452 | ✗ | device.SharedDevice() ? device::BufferUsage::Interop : device::BufferUsage::None; | |
| 453 | ✗ | Ref<device::Buffer> buffer = | |
| 454 | ✗ | device.CreateBuffer({.usage = usage, .label = "colorized"}, pixels.bytes); | |
| 455 | ✗ | return device::Tensor(std::move(buffer), 0, pixels.type, | |
| 456 | ✗ | {pixels.rows, pixels.cols, pixels.channels}); | |
| 457 | ✗ | } | |
| 458 | |||
| 459 | 30 | Pixels Packed(const Pixels& like, const std::vector<uint32_t>& rgba) { | |
| 460 | 30 | Pixels packed; | |
| 461 | 30 | packed.rows = like.rows; | |
| 462 | 30 | packed.cols = like.cols; | |
| 463 | 30 | packed.channels = 4; | |
| 464 | 30 | packed.bytes.resize(rgba.size() * sizeof(uint32_t)); | |
| 465 | 30 | std::memcpy(packed.bytes.data(), rgba.data(), packed.bytes.size()); | |
| 466 | 30 | return packed; | |
| 467 | ✗ | } | |
| 468 | |||
| 469 | 41 | void RequireScalars(device::DType type, int64_t channels, const ColorizeOptions& options) { | |
| 470 | 41 | if (type != device::DType::Float32) { | |
| 471 | 2 | throw Error("Scalars are colorized from float32 (rows, cols[, channels]) pixels"); | |
| 472 | } | ||
| 473 | 40 | if (int64_t(options.channel) >= channels) { | |
| 474 | 2 | throw Error("Colorizing channel {} of pixels with only {} channels", options.channel, | |
| 475 | 1 | channels); | |
| 476 | } | ||
| 477 | 39 | if (options.logarithmic && options.centered) { | |
| 478 | 2 | throw Error("Scalars are colorized logarithmically or centered on 0, not both"); | |
| 479 | } | ||
| 480 | 38 | } | |
| 481 | |||
| 482 | 61 | uint32_t ScalarFlags(const ColorizeOptions& options) { | |
| 483 | 110 | return (options.minimum ? kFixedMinimum : 0) | (options.maximum ? kFixedMaximum : 0) | | |
| 484 | 95 | (options.logarithmic ? kLogarithmic : 0) | (options.reversed ? kReversed : 0) | | |
| 485 | 61 | (options.centered ? kCentered : 0); | |
| 486 | } | ||
| 487 | |||
| 488 | 7 | void RequireHdr(device::DType type, int64_t channels) { | |
| 489 | 7 | if (type != device::DType::Float32 || (channels != 3 && channels != 4)) { | |
| 490 | 2 | throw Error("HDR color is colorized from float32 (rows, cols, 3 or 4) pixels"); | |
| 491 | } | ||
| 492 | 6 | } | |
| 493 | |||
| 494 | 5 | void RequireNormals(device::DType type, int64_t channels) { | |
| 495 | 5 | if (type != device::DType::Float32 || (channels != 3 && channels != 4)) { | |
| 496 | 2 | throw Error("Normals are colorized from float32 (rows, cols, 3 or 4) pixels"); | |
| 497 | } | ||
| 498 | 4 | } | |
| 499 | |||
| 500 | 5 | void RequireIds(device::DType type, int64_t channels) { | |
| 501 | 5 | if (type != device::DType::UInt32 || (channels != 1 && channels != 4)) { | |
| 502 | 2 | throw Error("Ids are colorized from uint32 (rows, cols[, 1 or 4]) pixels"); | |
| 503 | } | ||
| 504 | 4 | } | |
| 505 | |||
| 506 | 63304 | uint32_t PixelId(Span<const uint32_t> words) { | |
| 507 | 63304 | uint32_t key = words.empty() ? 0 : words[0]; | |
| 508 | 63304 | bool empty = key == 0; | |
| 509 | 253204 | for (size_t c = 1; c < words.size(); ++c) { | |
| 510 | 189900 | key = MixIdWord(key, words[c]); | |
| 511 | 191135 | empty = empty && words[c] == 0; | |
| 512 | } | ||
| 513 | 63304 | return empty ? 0 : IdRgba(key); | |
| 514 | } | ||
| 515 | |||
| 516 | struct Layout { | ||
| 517 | int64_t rows = 0; | ||
| 518 | int64_t cols = 0; | ||
| 519 | int64_t channels = 1; | ||
| 520 | bool compact = false; | ||
| 521 | }; | ||
| 522 | |||
| 523 | 50 | Layout Describe(const device::Tensor& tensor) { | |
| 524 | 50 | const Span<const int64_t> shape = tensor.Shape(); | |
| 525 | 50 | if ((shape.size() != 2 && shape.size() != 3) || shape[0] < 1 || shape[1] < 1) { | |
| 526 | ✗ | throw Error("Pixels come from (rows, cols[, channels]) tensors, not {} dimensions", | |
| 527 | ✗ | shape.size()); | |
| 528 | } | ||
| 529 | 50 | Layout layout{shape[0], shape[1], shape.size() == 3 ? shape[2] : 1, true}; | |
| 530 | 50 | const Span<const int64_t> strides = tensor.Strides(); | |
| 531 | 50 | int64_t expected = 1; | |
| 532 | 166 | for (size_t axis = shape.size(); axis-- > 0;) { | |
| 533 | 116 | layout.compact = layout.compact && (shape[axis] <= 1 || strides[axis] == expected); | |
| 534 | 116 | expected *= shape[axis]; | |
| 535 | } | ||
| 536 | 50 | return layout; | |
| 537 | } | ||
| 538 | |||
| 539 | 22 | device::Tensor Launch(const std::vector<const char*>& entries, const device::Tensor& input, | |
| 540 | ColorizeParams params, const Ref<device::Buffer>& range) { | ||
| 541 | 22 | const Layout layout = Describe(input); | |
| 542 | 22 | const uint64_t count = uint64_t(layout.rows) * uint64_t(layout.cols); | |
| 543 | 22 | if (count > std::numeric_limits<uint32_t>::max()) { | |
| 544 | ✗ | throw Error("Colorizing {} pixels; it takes at most {}", count, | |
| 545 | ✗ | std::numeric_limits<uint32_t>::max()); | |
| 546 | } | ||
| 547 | 22 | device::Device& device = input.GetDevice(); | |
| 548 | 22 | const Ref<device::Program> program = | |
| 549 | 22 | device.CreateProgram({.modules = {"chrono_batch_render.io.kernels.colorize"}, | |
| 550 | .library = &kIoKernels, | ||
| 551 | 22 | .entries = {entries.begin(), entries.end()}, | |
| 552 | 22 | .label = "colorize"}); | |
| 553 | 22 | device::Tensor output = | |
| 554 | 44 | device.CreateTensor(device::DType::UInt8, {layout.rows, layout.cols, 4}, "colorized"); | |
| 555 | 22 | const uint64_t words = count * uint64_t(layout.channels); | |
| 556 | 22 | params.input = device::RequireRange<uint32_t>( | |
| 557 | 22 | *program, "Colorize", input.GetBuffer().Slice(input.Offset(), words * 4), words); | |
| 558 | 22 | params.output = {output.Address()}; | |
| 559 | 22 | params.count = uint32_t(count); | |
| 560 | 22 | params.channels = uint32_t(layout.channels); | |
| 561 | |||
| 562 | 22 | device::Queue& queue = device.GetQueue(); | |
| 563 | 22 | device::CommandEncoder encoder = queue.CreateEncoder(); | |
| 564 | 22 | const device::Size3 groups = device::FoldGroups(device::GroupCount(count, kColorizeThreads)); | |
| 565 | 56 | for (const char* entry : entries) { | |
| 566 | 34 | encoder.Barrier(); | |
| 567 | 68 | encoder.Dispatch(program->Entry(entry), params, groups); | |
| 568 | } | ||
| 569 | 22 | const device::Operation done = queue.Submit(encoder.Finish()); | |
| 570 | 66 | queue.Retire(std::make_shared< | |
| 571 | 44 | std::tuple<Ref<device::Program>, Ref<device::Buffer>, Ref<device::Buffer>>>( | |
| 572 | 22 | program, Ref<device::Buffer>(&input.GetBuffer()), range), | |
| 573 | done); | ||
| 574 | 22 | return output; | |
| 575 | 88 | } | |
| 576 | |||
| 577 | 23 | ScalarBounds HostBounds(const Pixels& values, const ColorizeOptions& options) { | |
| 578 | 23 | const Span<const float> all = values.Values<float>(); | |
| 579 | 23 | const uint32_t flags = ScalarFlags(options); | |
| 580 | 23 | const size_t channels = size_t(values.channels); | |
| 581 | 23 | const size_t count = size_t(values.rows) * size_t(values.cols); | |
| 582 | 23 | float low = std::numeric_limits<float>::infinity(); | |
| 583 | 23 | float high = -std::numeric_limits<float>::infinity(); | |
| 584 | 511107 | for (size_t i = 0; i < count; ++i) { | |
| 585 | 511084 | const float value = all[i * channels + options.channel]; | |
| 586 | 997518 | if (ScalarShown(value, flags)) { | |
| 587 | 486434 | low = std::min(low, value); | |
| 588 | 493298 | high = std::max(high, value); | |
| 589 | } | ||
| 590 | } | ||
| 591 | 27 | return ResolveBounds(low, high, options.minimum.value_or(0), options.maximum.value_or(0), | |
| 592 | 23 | flags); | |
| 593 | } | ||
| 594 | |||
| 595 | } | ||
| 596 | |||
| 597 | 4 | std::array<uint8_t, 4> ColormapColor(Colormap colormap, float position) { | |
| 598 | 4 | const uint32_t rgba = ColormapRgba(uint32_t(colormap), position); | |
| 599 | 4 | return {uint8_t(rgba), uint8_t(rgba >> 8), uint8_t(rgba >> 16), uint8_t(rgba >> 24)}; | |
| 600 | } | ||
| 601 | |||
| 602 | 23 | Pixels ColorizeScalars(const Pixels& values, const ColorizeOptions& options) { | |
| 603 | 23 | RequireScalars(values.type, values.channels, options); | |
| 604 | 23 | const Span<const float> all = values.Values<float>(); | |
| 605 | 23 | const uint32_t flags = ScalarFlags(options); | |
| 606 | 23 | const size_t channels = size_t(values.channels); | |
| 607 | 23 | const size_t count = size_t(values.rows) * size_t(values.cols); | |
| 608 | 23 | const ScalarBounds bounds = HostBounds(values, options); | |
| 609 | 23 | std::vector<uint32_t> rgba(count); | |
| 610 | 511107 | for (size_t i = 0; i < count; ++i) { | |
| 611 | 511084 | rgba[i] = ScalarRgba(all[i * channels + options.channel], bounds, flags, | |
| 612 | 511084 | uint32_t(options.colormap)); | |
| 613 | } | ||
| 614 | 23 | return Packed(values, rgba); | |
| 615 | 23 | } | |
| 616 | |||
| 617 | 4 | Pixels ColorizeHdr(const Pixels& color, float exposure) { | |
| 618 | 4 | RequireHdr(color.type, color.channels); | |
| 619 | 4 | const Span<const float> values = color.Values<float>(); | |
| 620 | 4 | const float scale = std::exp2(exposure); | |
| 621 | 4 | std::vector<uint32_t> rgba(size_t(color.rows) * size_t(color.cols)); | |
| 622 | 63309 | for (size_t i = 0; i < rgba.size(); ++i) { | |
| 623 | 63305 | const float* pixel = values.data() + i * size_t(color.channels); | |
| 624 | 63305 | rgba[i] = | |
| 625 | 63305 | HdrRgba(pixel[0], pixel[1], pixel[2], color.channels == 4 ? pixel[3] : 1.0f, scale); | |
| 626 | } | ||
| 627 | 4 | return Packed(color, rgba); | |
| 628 | 4 | } | |
| 629 | |||
| 630 | 1 | Pixels ColorizeNormals(const Pixels& normals) { | |
| 631 | 1 | RequireNormals(normals.type, normals.channels); | |
| 632 | 1 | const Span<const float> values = normals.Values<float>(); | |
| 633 | 1 | std::vector<uint32_t> rgba(size_t(normals.rows) * size_t(normals.cols)); | |
| 634 | 63301 | for (size_t i = 0; i < rgba.size(); ++i) { | |
| 635 | 63300 | const float* normal = values.data() + i * size_t(normals.channels); | |
| 636 | 63300 | rgba[i] = NormalRgba(normal[0], normal[1], normal[2]); | |
| 637 | } | ||
| 638 | 1 | return Packed(normals, rgba); | |
| 639 | 1 | } | |
| 640 | |||
| 641 | ✗ | std::array<uint8_t, 4> IdColor(Span<const uint32_t> words) { | |
| 642 | ✗ | const uint32_t rgba = PixelId(words); | |
| 643 | ✗ | return {uint8_t(rgba), uint8_t(rgba >> 8), uint8_t(rgba >> 16), uint8_t(rgba >> 24)}; | |
| 644 | } | ||
| 645 | |||
| 646 | 2 | Pixels ColorizeIds(const Pixels& ids) { | |
| 647 | 2 | RequireIds(ids.type, ids.channels); | |
| 648 | 2 | const Span<const uint32_t> words = ids.Values<uint32_t>(); | |
| 649 | 2 | const size_t channels = size_t(ids.channels); | |
| 650 | 2 | std::vector<uint32_t> rgba(size_t(ids.rows) * size_t(ids.cols)); | |
| 651 | 63306 | for (size_t i = 0; i < rgba.size(); ++i) { | |
| 652 | 63304 | rgba[i] = PixelId(Span<const uint32_t>(words.data() + i * channels, channels)); | |
| 653 | } | ||
| 654 | 2 | return Packed(ids, rgba); | |
| 655 | 2 | } | |
| 656 | |||
| 657 | namespace { | ||
| 658 | |||
| 659 | 18 | device::Tensor ScalarTensor(const device::Tensor& values, const ColorizeOptions& options, | |
| 660 | device::Tensor* bounds) { | ||
| 661 | 18 | const Layout layout = Describe(values); | |
| 662 | 18 | RequireScalars(values.Type(), layout.channels, options); | |
| 663 | 15 | if (!layout.compact) { | |
| 664 | ✗ | const Pixels pixels(values); | |
| 665 | ✗ | if (bounds) { | |
| 666 | ✗ | const ScalarBounds used = HostBounds(pixels, options); | |
| 667 | ✗ | const float words[2] = {used.low, used.high}; | |
| 668 | ✗ | *bounds = device::Tensor(values.GetDevice().CreateBuffer({.label = "colorize bounds"}, | |
| 669 | Span<const float>(words)), | ||
| 670 | ✗ | 0, device::DType::Float32, {2}); | |
| 671 | } | ||
| 672 | ✗ | return ToTensor(values, ColorizeScalars(pixels, options)); | |
| 673 | ✗ | } | |
| 674 | 15 | ColorizeParams params{}; | |
| 675 | 15 | params.channel = options.channel; | |
| 676 | 15 | params.colormap = uint32_t(options.colormap); | |
| 677 | 15 | params.flags = ScalarFlags(options); | |
| 678 | 15 | params.minimum = options.minimum.value_or(0); | |
| 679 | 15 | params.maximum = options.maximum.value_or(0); | |
| 680 | 15 | if (bounds) { | |
| 681 | ✗ | *bounds = values.GetDevice().CreateTensor(device::DType::Float32, {2}, "colorize bounds"); | |
| 682 | ✗ | params.bounds = {bounds->Address()}; | |
| 683 | ✗ | params.flags |= kWriteBounds; | |
| 684 | } | ||
| 685 | 15 | std::vector<const char*> entries = {"ColorizeScalars"}; | |
| 686 | 15 | const uint32_t initial[2] = {kNoRange, 0}; | |
| 687 | 15 | const Ref<device::Buffer> range = | |
| 688 | 15 | values.GetDevice().CreateBuffer({.label = "colorize range"}, Span<const uint32_t>(initial)); | |
| 689 | 15 | params.range = {range->Address()}; | |
| 690 | 15 | if (!options.minimum || !options.maximum) { | |
| 691 | 12 | entries.insert(entries.begin(), "ScalarRange"); | |
| 692 | } | ||
| 693 | 15 | return Launch(entries, values, params, range); | |
| 694 | 45 | } | |
| 695 | |||
| 696 | } | ||
| 697 | |||
| 698 | 18 | device::Tensor ColorizeScalars(const device::Tensor& values, const ColorizeOptions& options) { | |
| 699 | 18 | return ScalarTensor(values, options, nullptr); | |
| 700 | } | ||
| 701 | |||
| 702 | ✗ | device::Tensor ColorizeScalars(const device::Tensor& values, const ColorizeOptions& options, | |
| 703 | device::Tensor& bounds) { | ||
| 704 | ✗ | return ScalarTensor(values, options, &bounds); | |
| 705 | } | ||
| 706 | |||
| 707 | 3 | device::Tensor ColorizeHdr(const device::Tensor& color, float exposure) { | |
| 708 | 3 | const Layout layout = Describe(color); | |
| 709 | 3 | RequireHdr(color.Type(), layout.channels); | |
| 710 | 2 | if (!layout.compact) { | |
| 711 | ✗ | return ToTensor(color, ColorizeHdr(Pixels(color), exposure)); | |
| 712 | } | ||
| 713 | 2 | ColorizeParams params{}; | |
| 714 | 2 | params.scale = std::exp2(exposure); | |
| 715 | 2 | return Launch({"ColorizeHdr"}, color, params, {}); | |
| 716 | } | ||
| 717 | |||
| 718 | 4 | device::Tensor ColorizeNormals(const device::Tensor& normals) { | |
| 719 | 4 | const Layout layout = Describe(normals); | |
| 720 | 4 | RequireNormals(normals.Type(), layout.channels); | |
| 721 | 3 | if (!layout.compact) { | |
| 722 | ✗ | return ToTensor(normals, ColorizeNormals(Pixels(normals))); | |
| 723 | } | ||
| 724 | 3 | return Launch({"ColorizeNormals"}, normals, {}, {}); | |
| 725 | } | ||
| 726 | |||
| 727 | 3 | device::Tensor ColorizeIds(const device::Tensor& ids) { | |
| 728 | 3 | const Layout layout = Describe(ids); | |
| 729 | 3 | RequireIds(ids.Type(), layout.channels); | |
| 730 | 2 | if (!layout.compact) { | |
| 731 | ✗ | return ToTensor(ids, ColorizeIds(Pixels(ids))); | |
| 732 | } | ||
| 733 | 2 | return Launch({"ColorizeIds"}, ids, {}, {}); | |
| 734 | } | ||
| 735 | |||
| 736 | } | ||
| 737 |