lib/extras/enc/exr.cc

   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 = 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 = nullptr) 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