1 files changed, 18 insertions, 10 deletions
diff --git a/tools/hdr/tone_map.cc b/tools/hdr/tone_map.cc
index 1ef3823..67fea48 100644
--- a/tools/hdr/tone_map.cc
+++ b/tools/hdr/tone_map.cc
@@ -8,13 +8,14 @@
 
 #include "lib/extras/codec.h"
 #include "lib/extras/tone_mapping.h"
-#include "lib/jxl/base/thread_pool_internal.h"
-#include "lib/jxl/enc_color_management.h"
 #include "tools/args.h"
 #include "tools/cmdline.h"
+#include "tools/file_io.h"
+#include "tools/hdr/image_utils.h"
+#include "tools/thread_pool_internal.h"
 
 int main(int argc, const char** argv) {
-  jxl::ThreadPoolInternal pool;
+  jpegxl::tools::ThreadPoolInternal pool;
 
   jpegxl::tools::CommandLineParser parser;
   float max_nits = 0;
@@ -69,7 +70,9 @@ int main(int argc, const char** argv) {
   jxl::CodecInOut image;
   jxl::extras::ColorHints color_hints;
   color_hints.Add("color_space", "RGB_D65_202_Rel_PeQ");
-  JXL_CHECK(jxl::SetFromFile(input_filename, color_hints, &image, &pool));
+  std::vector<uint8_t> encoded;
+  JXL_CHECK(jpegxl::tools::ReadFile(input_filename, &encoded));
+  JXL_CHECK(jxl::SetFromBytes(jxl::Bytes(encoded), color_hints, &image, &pool));
   if (max_nits > 0) {
     image.metadata.m.SetIntensityTarget(max_nits);
   }
@@ -77,13 +80,18 @@ int main(int argc, const char** argv) {
   JXL_CHECK(jxl::GamutMap(&image, preserve_saturation, &pool));
 
   jxl::ColorEncoding c_out = image.metadata.m.color_encoding;
-  if (pq) {
-    c_out.tf.SetTransferFunction(jxl::TransferFunction::kPQ);
-  } else {
-    c_out.tf.SetTransferFunction(jxl::TransferFunction::k709);
+  jxl::cms::TransferFunction tf =
+      pq ? jxl::TransferFunction::kPQ : jxl::TransferFunction::kSRGB;
+
+  if (jxl::extras::CodecFromPath(output_filename) == jxl::extras::Codec::kEXR) {
+    tf = jxl::TransferFunction::kLinear;
+    image.metadata.m.SetFloat16Samples();
   }
+  c_out.Tf().SetTransferFunction(tf);
+
   JXL_CHECK(c_out.CreateICC());
-  JXL_CHECK(image.TransformTo(c_out, jxl::GetJxlCms(), &pool));
+  JXL_CHECK(jpegxl::tools::TransformCodecInOutTo(image, c_out, &pool));
   image.metadata.m.color_encoding = c_out;
-  JXL_CHECK(jxl::EncodeToFile(image, output_filename, &pool));
+  JXL_CHECK(jxl::Encode(image, output_filename, &encoded, &pool));
+  JXL_CHECK(jpegxl::tools::WriteFile(output_filename, encoded));
 }