Source code
Revision control
Copy as Markdown
Other Tools
// Copyright (c) the JPEG XL Project Authors. All rights reserved.
//
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
#include "lib/extras/enc/exr.h"
#if JPEGXL_ENABLE_EXR
#include <ImfChromaticitiesAttribute.h>
#include <ImfIO.h>
#include <ImfRgbaFile.h>
#include <ImfStandardAttributes.h>
#endif
#include <jxl/codestream_header.h>
#include <vector>
#include "lib/extras/packed_image.h"
#include "lib/jxl/base/byte_order.h"
namespace jxl {
namespace extras {
#if JPEGXL_ENABLE_EXR
namespace {
namespace OpenEXR = OPENEXR_IMF_NAMESPACE;
namespace Imath = IMATH_NAMESPACE;
// OpenEXR::Int64 is deprecated in favor of using uint64_t directly, but using
// uint64_t as recommended causes build failures with previous OpenEXR versions
// on macOS, where the definition for OpenEXR::Int64 was actually not equivalent
// to uint64_t. This alternative should work in all cases.
using ExrInt64 = decltype(std::declval<OpenEXR::IStream>().tellg());
class InMemoryOStream : public OpenEXR::OStream {
public:
// `bytes` must outlive the InMemoryOStream.
explicit InMemoryOStream(std::vector<uint8_t>* const bytes)
: OStream(/*fileName=*/""), bytes_(*bytes) {}
void write(const char c[], const int n) override {
if (bytes_.size() < pos_ + n) {
bytes_.resize(pos_ + n);
}
std::copy_n(c, n, bytes_.begin() + pos_);
pos_ += n;
}
ExrInt64 tellp() override { return pos_; }
void seekp(const ExrInt64 pos) override {
if (bytes_.size() + 1 < pos) {
bytes_.resize(pos - 1);
}
pos_ = pos;
}
private:
std::vector<uint8_t>& bytes_;
size_t pos_ = 0;
};
// Loads a Big-Endian float
float LoadBEFloat(const uint8_t* p) {
uint32_t u = LoadBE32(p);
float result;
memcpy(&result, &u, 4);
return result;
}
// Loads a Little-Endian float
float LoadLEFloat(const uint8_t* p) {
uint32_t u = LoadLE32(p);
float result;
memcpy(&result, &u, 4);
return result;
}
Status EncodeImageEXR(const PackedImage& image, const JxlBasicInfo& info,
const JxlColorEncoding& c_enc, ThreadPool* pool,
std::vector<uint8_t>* bytes) {
OpenEXR::setGlobalThreadCount(0);
const size_t xsize = info.xsize;
const size_t ysize = info.ysize;
const bool has_alpha = info.alpha_bits > 0;
const bool alpha_is_premultiplied = FROM_JXL_BOOL(info.alpha_premultiplied);
if (info.num_color_channels != 3 ||
c_enc.color_space != JXL_COLOR_SPACE_RGB ||
c_enc.transfer_function != JXL_TRANSFER_FUNCTION_LINEAR) {
return JXL_FAILURE("Unsupported color encoding for OpenEXR output.");
}
const size_t num_channels = 3 + (has_alpha ? 1 : 0);
const JxlPixelFormat format = image.format;
if (format.data_type != JXL_TYPE_FLOAT) {
return JXL_FAILURE("Unsupported pixel format for OpenEXR output");
}
const uint8_t* in = reinterpret_cast<const uint8_t*>(image.pixels());
size_t in_stride = num_channels * 4 * xsize;
OpenEXR::Header header(xsize, ysize);
OpenEXR::Chromaticities chromaticities;
chromaticities.red =
Imath::V2f(c_enc.primaries_red_xy[0], c_enc.primaries_red_xy[1]);
chromaticities.green =
Imath::V2f(c_enc.primaries_green_xy[0], c_enc.primaries_green_xy[1]);
chromaticities.blue =
Imath::V2f(c_enc.primaries_blue_xy[0], c_enc.primaries_blue_xy[1]);
chromaticities.white =
Imath::V2f(c_enc.white_point_xy[0], c_enc.white_point_xy[1]);
OpenEXR::addChromaticities(header, chromaticities);
OpenEXR::addWhiteLuminance(header, info.intensity_target);
auto loadFloat =
format.endianness == JXL_BIG_ENDIAN ? LoadBEFloat : LoadLEFloat;
auto loadAlpha =
has_alpha ? loadFloat : [](const uint8_t* p) -> float { return 1.0f; };
// Ensure that the destructor of RgbaOutputFile has run before we look at the
// size of `bytes`.
{
InMemoryOStream os(bytes);
OpenEXR::RgbaOutputFile output(
os, header, has_alpha ? OpenEXR::WRITE_RGBA : OpenEXR::WRITE_RGB);
// How many rows to write at once. Again, the OpenEXR documentation
// recommends writing the whole image in one call.
const int y_chunk_size = ysize;
std::vector<OpenEXR::Rgba> output_rows(xsize * y_chunk_size);
for (size_t start_y = 0; start_y < ysize; start_y += y_chunk_size) {
// Inclusive.
const size_t end_y = std::min(start_y + y_chunk_size - 1, ysize - 1);
output.setFrameBuffer(output_rows.data() - start_y * xsize,
/*xStride=*/1, /*yStride=*/xsize);
for (size_t y = start_y; y <= end_y; ++y) {
const uint8_t* in_row = &in[(y - start_y) * in_stride];
OpenEXR::Rgba* const JXL_RESTRICT row_data =
&output_rows[(y - start_y) * xsize];
for (size_t x = 0; x < xsize; ++x) {
const uint8_t* in_pixel = &in_row[4 * num_channels * x];
float r = loadFloat(&in_pixel[0]);
float g = loadFloat(&in_pixel[4]);
float b = loadFloat(&in_pixel[8]);
const float alpha = loadAlpha(&in_pixel[12]);
if (!alpha_is_premultiplied) {
r *= alpha;
g *= alpha;
b *= alpha;
}
row_data[x] = OpenEXR::Rgba(r, g, b, alpha);
}
}
output.writePixels(/*numScanLines=*/end_y - start_y + 1);
}
}
return true;
}
class EXREncoder : public Encoder {
std::vector<JxlPixelFormat> AcceptedFormats() const override {
std::vector<JxlPixelFormat> formats;
for (const uint32_t num_channels : {1, 2, 3, 4}) {
for (const JxlDataType data_type : {JXL_TYPE_FLOAT}) {
for (JxlEndianness endianness : {JXL_BIG_ENDIAN, JXL_LITTLE_ENDIAN}) {
formats.push_back(JxlPixelFormat{/*num_channels=*/num_channels,
/*data_type=*/data_type,
/*endianness=*/endianness,
/*align=*/0});
}
}
}
return formats;
}
Status Encode(const PackedPixelFile& ppf, EncodedImage* encoded_image,
ThreadPool* pool) const override {
JXL_RETURN_IF_ERROR(VerifyBasicInfo(ppf.info));
encoded_image->icc.clear();
encoded_image->bitstreams.clear();
encoded_image->bitstreams.reserve(ppf.frames.size());
for (const auto& frame : ppf.frames) {
JXL_RETURN_IF_ERROR(VerifyPackedImage(frame.color, ppf.info));
encoded_image->bitstreams.emplace_back();
JXL_RETURN_IF_ERROR(EncodeImageEXR(frame.color, ppf.info,
ppf.color_encoding, pool,
&encoded_image->bitstreams.back()));
}
return true;
}
};
} // namespace
#endif
std::unique_ptr<Encoder> GetEXREncoder() {
#if JPEGXL_ENABLE_EXR
return jxl::make_unique<EXREncoder>();
#else
return nullptr;
#endif
}
} // namespace extras
} // namespace jxl