exr.cc (6910B)
1 // Copyright (c) the JPEG XL Project Authors. All rights reserved. 2 // 3 // Use of this source code is governed by a BSD-style 4 // license that can be found in the LICENSE file. 5 6 #include "lib/extras/enc/exr.h" 7 8 #if JPEGXL_ENABLE_EXR 9 #include <ImfChromaticitiesAttribute.h> 10 #include <ImfIO.h> 11 #include <ImfRgbaFile.h> 12 #include <ImfStandardAttributes.h> 13 #endif 14 #include <jxl/codestream_header.h> 15 16 #include <vector> 17 18 #include "lib/extras/packed_image.h" 19 #include "lib/jxl/base/byte_order.h" 20 21 namespace jxl { 22 namespace extras { 23 24 #if JPEGXL_ENABLE_EXR 25 namespace { 26 27 namespace OpenEXR = OPENEXR_IMF_NAMESPACE; 28 namespace Imath = IMATH_NAMESPACE; 29 30 // OpenEXR::Int64 is deprecated in favor of using uint64_t directly, but using 31 // uint64_t as recommended causes build failures with previous OpenEXR versions 32 // on macOS, where the definition for OpenEXR::Int64 was actually not equivalent 33 // to uint64_t. This alternative should work in all cases. 34 using ExrInt64 = decltype(std::declval<OpenEXR::IStream>().tellg()); 35 36 class InMemoryOStream : public OpenEXR::OStream { 37 public: 38 // `bytes` must outlive the InMemoryOStream. 39 explicit InMemoryOStream(std::vector<uint8_t>* const bytes) 40 : OStream(/*fileName=*/""), bytes_(*bytes) {} 41 42 void write(const char c[], const int n) override { 43 if (bytes_.size() < pos_ + n) { 44 bytes_.resize(pos_ + n); 45 } 46 std::copy_n(c, n, bytes_.begin() + pos_); 47 pos_ += n; 48 } 49 50 ExrInt64 tellp() override { return pos_; } 51 void seekp(const ExrInt64 pos) override { 52 if (bytes_.size() + 1 < pos) { 53 bytes_.resize(pos - 1); 54 } 55 pos_ = pos; 56 } 57 58 private: 59 std::vector<uint8_t>& bytes_; 60 size_t pos_ = 0; 61 }; 62 63 // Loads a Big-Endian float 64 float LoadBEFloat(const uint8_t* p) { 65 uint32_t u = LoadBE32(p); 66 float result; 67 memcpy(&result, &u, 4); 68 return result; 69 } 70 71 // Loads a Little-Endian float 72 float LoadLEFloat(const uint8_t* p) { 73 uint32_t u = LoadLE32(p); 74 float result; 75 memcpy(&result, &u, 4); 76 return result; 77 } 78 79 Status EncodeImageEXR(const PackedImage& image, const JxlBasicInfo& info, 80 const JxlColorEncoding& c_enc, ThreadPool* pool, 81 std::vector<uint8_t>* bytes) { 82 OpenEXR::setGlobalThreadCount(0); 83 84 const size_t xsize = info.xsize; 85 const size_t ysize = info.ysize; 86 const bool has_alpha = info.alpha_bits > 0; 87 const bool alpha_is_premultiplied = FROM_JXL_BOOL(info.alpha_premultiplied); 88 89 if (info.num_color_channels != 3 || 90 c_enc.color_space != JXL_COLOR_SPACE_RGB || 91 c_enc.transfer_function != JXL_TRANSFER_FUNCTION_LINEAR) { 92 return JXL_FAILURE("Unsupported color encoding for OpenEXR output."); 93 } 94 95 const size_t num_channels = 3 + (has_alpha ? 1 : 0); 96 const JxlPixelFormat format = image.format; 97 98 if (format.data_type != JXL_TYPE_FLOAT) { 99 return JXL_FAILURE("Unsupported pixel format for OpenEXR output"); 100 } 101 102 const uint8_t* in = reinterpret_cast<const uint8_t*>(image.pixels()); 103 size_t in_stride = num_channels * 4 * xsize; 104 105 OpenEXR::Header header(xsize, ysize); 106 OpenEXR::Chromaticities chromaticities; 107 chromaticities.red = 108 Imath::V2f(c_enc.primaries_red_xy[0], c_enc.primaries_red_xy[1]); 109 chromaticities.green = 110 Imath::V2f(c_enc.primaries_green_xy[0], c_enc.primaries_green_xy[1]); 111 chromaticities.blue = 112 Imath::V2f(c_enc.primaries_blue_xy[0], c_enc.primaries_blue_xy[1]); 113 chromaticities.white = 114 Imath::V2f(c_enc.white_point_xy[0], c_enc.white_point_xy[1]); 115 OpenEXR::addChromaticities(header, chromaticities); 116 OpenEXR::addWhiteLuminance(header, info.intensity_target); 117 118 auto loadFloat = 119 format.endianness == JXL_BIG_ENDIAN ? LoadBEFloat : LoadLEFloat; 120 auto loadAlpha = 121 has_alpha ? loadFloat : [](const uint8_t* p) -> float { return 1.0f; }; 122 123 // Ensure that the destructor of RgbaOutputFile has run before we look at the 124 // size of `bytes`. 125 { 126 InMemoryOStream os(bytes); 127 OpenEXR::RgbaOutputFile output( 128 os, header, has_alpha ? OpenEXR::WRITE_RGBA : OpenEXR::WRITE_RGB); 129 // How many rows to write at once. Again, the OpenEXR documentation 130 // recommends writing the whole image in one call. 131 const int y_chunk_size = ysize; 132 std::vector<OpenEXR::Rgba> output_rows(xsize * y_chunk_size); 133 134 for (size_t start_y = 0; start_y < ysize; start_y += y_chunk_size) { 135 // Inclusive. 136 const size_t end_y = std::min(start_y + y_chunk_size - 1, ysize - 1); 137 output.setFrameBuffer(output_rows.data() - start_y * xsize, 138 /*xStride=*/1, /*yStride=*/xsize); 139 for (size_t y = start_y; y <= end_y; ++y) { 140 const uint8_t* in_row = &in[(y - start_y) * in_stride]; 141 OpenEXR::Rgba* const JXL_RESTRICT row_data = 142 &output_rows[(y - start_y) * xsize]; 143 for (size_t x = 0; x < xsize; ++x) { 144 const uint8_t* in_pixel = &in_row[4 * num_channels * x]; 145 float r = loadFloat(&in_pixel[0]); 146 float g = loadFloat(&in_pixel[4]); 147 float b = loadFloat(&in_pixel[8]); 148 const float alpha = loadAlpha(&in_pixel[12]); 149 if (!alpha_is_premultiplied) { 150 r *= alpha; 151 g *= alpha; 152 b *= alpha; 153 } 154 row_data[x] = OpenEXR::Rgba(r, g, b, alpha); 155 } 156 } 157 output.writePixels(/*numScanLines=*/end_y - start_y + 1); 158 } 159 } 160 161 return true; 162 } 163 164 class EXREncoder : public Encoder { 165 std::vector<JxlPixelFormat> AcceptedFormats() const override { 166 std::vector<JxlPixelFormat> formats; 167 for (const uint32_t num_channels : {1, 2, 3, 4}) { 168 for (const JxlDataType data_type : {JXL_TYPE_FLOAT}) { 169 for (JxlEndianness endianness : {JXL_BIG_ENDIAN, JXL_LITTLE_ENDIAN}) { 170 formats.push_back(JxlPixelFormat{/*num_channels=*/num_channels, 171 /*data_type=*/data_type, 172 /*endianness=*/endianness, 173 /*align=*/0}); 174 } 175 } 176 } 177 return formats; 178 } 179 Status Encode(const PackedPixelFile& ppf, EncodedImage* encoded_image, 180 ThreadPool* pool) const override { 181 JXL_RETURN_IF_ERROR(VerifyBasicInfo(ppf.info)); 182 encoded_image->icc.clear(); 183 encoded_image->bitstreams.clear(); 184 encoded_image->bitstreams.reserve(ppf.frames.size()); 185 for (const auto& frame : ppf.frames) { 186 JXL_RETURN_IF_ERROR(VerifyPackedImage(frame.color, ppf.info)); 187 encoded_image->bitstreams.emplace_back(); 188 JXL_RETURN_IF_ERROR(EncodeImageEXR(frame.color, ppf.info, 189 ppf.color_encoding, pool, 190 &encoded_image->bitstreams.back())); 191 } 192 return true; 193 } 194 }; 195 196 } // namespace 197 #endif 198 199 std::unique_ptr<Encoder> GetEXREncoder() { 200 #if JPEGXL_ENABLE_EXR 201 return jxl::make_unique<EXREncoder>(); 202 #else 203 return nullptr; 204 #endif 205 } 206 207 } // namespace extras 208 } // namespace jxl