Merge pull request #1415 from zenustech/aov-save-exr

Aov save exr

ui/zenoedit/dialog/zrecorddlg.cpp

@@ -62,6 +62,7 @@ bool ZRecordVideoDlg::getInfo(VideoRecInfo &info)
 {
     auto &ud = zeno::getSession().userData();
     ud.set2("output_aov", m_ui->cbAOV->checkState() == Qt::Checked);
+    ud.set2("output_exr", m_ui->cbExportEXR->checkState() == Qt::Checked);
     auto &path = info.record_path;
     auto &fn = info.videoname;
     info.fps = m_ui->fps->text().toInt();
@@ -72,6 +73,7 @@ bool ZRecordVideoDlg::getInfo(VideoRecInfo &info)
     info.res[1] = m_ui->lineHeight->text().toFloat();
     path = m_ui->linePath->text();
     info.bExportVideo = m_ui->cbExportVideo->checkState() == Qt::Checked;
+    info.bExportEXR = m_ui->cbExportEXR->checkState() == Qt::Checked;
     info.needDenoise = m_ui->cbNeedDenoise->checkState() == Qt::Checked;
     info.bAutoRemoveCache = m_ui->cbRemoveAfterRender->checkState() == Qt::Checked;
     if (path.isEmpty())

ui/zenoedit/dialog/zrecorddlg.ui

@@ -130,6 +130,16 @@
      </item>
     </layout>
    </item>
+   <item>
+    <widget class="QCheckBox" name="cbExportEXR">
+     <property name="text">
+      <string>export exr</string>
+     </property>
+     <property name="checked">
+      <bool>false</bool>
+     </property>
+    </widget>
+   </item>
    <item>
     <widget class="QCheckBox" name="cbExportVideo">
      <property name="text">

ui/zenoedit/recordmain.cpp

@@ -77,6 +77,7 @@ int record_main(const QCoreApplication& app)
         {"optix", "optix", "optix mode"},
         {"video", "video", "export video"},
         {"aov", "aov", "aov"},
+        {"exr", "exr", "exr"},
         {"needDenoise", "needDenoise", "needDenoise"},
         {"videoname", "videoname", "export video's name"},
         {"subzsg", "subgraphzsg", "subgraph zsg file path"},
@@ -147,8 +148,10 @@ int record_main(const QCoreApplication& app)
 	param.isExportVideo = cmdParser.isSet("video") ? cmdParser.value("video").toInt() : 0;
 	param.needDenoise = cmdParser.isSet("needDenoise") ? cmdParser.value("needDenoise").toInt() : 0;
 	int enableAOV = cmdParser.isSet("aov") ? cmdParser.value("aov").toInt() : 0;
+	param.export_exr = cmdParser.isSet("exr") && cmdParser.value("exr").toInt() != 0;
     auto &ud = zeno::getSession().userData();
     ud.set2("output_aov", enableAOV != 0);
+    ud.set2("output_exr", param.export_exr);
 	param.videoName = cmdParser.isSet("videoname") ? cmdParser.value("videoname") : "output.mp4";
 	param.subZsg = cmdParser.isSet("subzsg") ? cmdParser.value("subzsg") : "";
 #else

ui/zenoedit/viewport/recordvideomgr.cpp

@@ -94,6 +94,10 @@ void RecordVideoMgr::setRecordInfo(const VideoRecInfo& recInfo)
 
 void RecordVideoMgr::endRecToExportVideo()
 {
+    if (m_recordInfo.bExportEXR) {
+        emit recordFinished(m_recordInfo.record_path);
+        return;
+    }
     // denoising
     if (m_recordInfo.needDenoise) {
         QString dir_path = m_recordInfo.record_path + "/P/";

ui/zenoedit/viewport/recordvideomgr.h

@@ -20,6 +20,7 @@ struct VideoRecInfo
     bool exitWhenRecordFinish = false;
     bool bRecordByCommandLine = false;
     bool bAutoRemoveCache = false;
+    bool bExportEXR = false;
 };
 Q_DECLARE_METATYPE(VideoRecInfo);
 

ui/zenoedit/zenomainwindow.cpp

@@ -827,6 +827,7 @@ void ZenoMainWindow::optixRunRender(const ZENO_RECORD_RUN_INITPARAM& param, LAUN
     recInfo.videoname = param.videoName;
     recInfo.bExportVideo = param.isExportVideo;
     recInfo.needDenoise = param.needDenoise;
+    recInfo.bExportEXR = param.export_exr;
     recInfo.exitWhenRecordFinish = param.exitWhenRecordFinish;
 
     if (!param.sPixel.isEmpty())
@@ -918,6 +919,7 @@ void ZenoMainWindow::solidRunRender(const ZENO_RECORD_RUN_INITPARAM& param)
     recInfo.videoname = param.videoName;
     recInfo.bExportVideo = param.isExportVideo;
     recInfo.needDenoise = param.needDenoise;
+    recInfo.bExportEXR = param.export_exr;
     recInfo.exitWhenRecordFinish = param.exitWhenRecordFinish;
     recInfo.bRecordByCommandLine = true;
 

ui/zenoedit/zenomainwindow.h

@@ -18,6 +18,7 @@ struct ZENO_RECORD_RUN_INITPARAM {
     bool bOptix = false;    //is optix view.
     bool isExportVideo = false;
     bool needDenoise = false;
+    bool export_exr = false;
     int iFrame = 0;
     int iSFrame = 0;
     int iSample = 0;

zenovis/xinxinoptix/PTKernel.cu

@@ -299,10 +299,11 @@ extern "C" __global__ void __raygen__rg()
       float3 out_color_t = accum_color_t;
       float3 out_color_b = accum_color_b;
       params.frame_buffer[ image_index ] = make_color ( out_color );
-      params.frame_buffer_D[ image_index ] = make_color ( out_color_d );
-      params.frame_buffer_S[ image_index ] = make_color ( out_color_s );
-      params.frame_buffer_T[ image_index ] = make_color ( out_color_t );
-      params.frame_buffer_B[ image_index ] = make_color ( out_color_b );
+      params.frame_buffer_C[ image_index ] = accum_color;
+      params.frame_buffer_D[ image_index ] = accum_color_d;
+      params.frame_buffer_S[ image_index ] = accum_color_s;
+      params.frame_buffer_T[ image_index ] = accum_color_t;
+      params.frame_buffer_B[ image_index ] = accum_color_b;
 
       if (params.denoise) {
           params.albedo_buffer[ image_index ] = tmp_albedo;

zenovis/xinxinoptix/optixPathTracer.cpp

@@ -64,6 +64,7 @@
 #include <glm/gtc/type_ptr.hpp>
 
 #include "tinyexr.h"
+#include "zeno/utils/image_proc.h"
 
 #ifndef M_PI
 #define M_PI 3.14159265358979323846
@@ -154,10 +155,11 @@ typedef Record<HitGroupData> HitGroupRecord;
 //};
 
 std::optional<sutil::CUDAOutputBuffer<uchar4>> output_buffer_o;
-std::optional<sutil::CUDAOutputBuffer<uchar4>> output_buffer_diffuse;
-std::optional<sutil::CUDAOutputBuffer<uchar4>> output_buffer_specular;
-std::optional<sutil::CUDAOutputBuffer<uchar4>> output_buffer_transmit;
-std::optional<sutil::CUDAOutputBuffer<uchar4>> output_buffer_background;
+std::optional<sutil::CUDAOutputBuffer<float3>> output_buffer_color;
+std::optional<sutil::CUDAOutputBuffer<float3>> output_buffer_diffuse;
+std::optional<sutil::CUDAOutputBuffer<float3>> output_buffer_specular;
+std::optional<sutil::CUDAOutputBuffer<float3>> output_buffer_transmit;
+std::optional<sutil::CUDAOutputBuffer<float3>> output_buffer_background;
 using Vertex = float4;
 std::vector<Vertex> g_lightMesh;
 std::vector<Vertex> g_lightColor;
@@ -427,6 +429,7 @@ static void handleResize( sutil::CUDAOutputBuffer<uchar4>& output_buffer, Params
     resize_dirty = false;
 
     output_buffer.resize( params.width, params.height );
+    (*output_buffer_color).resize( params.width, params.height );
     (*output_buffer_diffuse).resize( params.width, params.height );
     (*output_buffer_specular).resize( params.width, params.height );
     (*output_buffer_transmit).resize( params.width, params.height );
@@ -492,6 +495,7 @@ static void launchSubframe( sutil::CUDAOutputBuffer<uchar4>& output_buffer, Path
     // Launch
     uchar4* result_buffer_data = output_buffer.map();
     state.params.frame_buffer  = result_buffer_data;
+    state.params.frame_buffer_C = (*output_buffer_color     ).map();
     state.params.frame_buffer_D = (*output_buffer_diffuse   ).map();
     state.params.frame_buffer_S = (*output_buffer_specular  ).map();
     state.params.frame_buffer_T = (*output_buffer_transmit  ).map();
@@ -528,6 +532,7 @@ static void launchSubframe( sutil::CUDAOutputBuffer<uchar4>& output_buffer, Path
       }
     }
     output_buffer.unmap();
+    (*output_buffer_color   ).unmap();
     (*output_buffer_diffuse   ).unmap();
     (*output_buffer_specular  ).unmap();
     (*output_buffer_transmit  ).unmap();
@@ -1608,6 +1613,14 @@ void optixinit( int argc, char* argv[] )
       );
       output_buffer_o->setStream( 0 );
     }
+    if (!output_buffer_color) {
+      output_buffer_color.emplace(
+          output_buffer_type,
+          state.params.width,
+          state.params.height
+      );
+      output_buffer_color->setStream( 0 );
+    }
     if (!output_buffer_diffuse) {
       output_buffer_diffuse.emplace(
           output_buffer_type,
@@ -3402,9 +3415,24 @@ void *optixgetimg_extra(std::string name) {
     else if (name == "background") {
         return output_buffer_background->getHostPointer();
     }
+    else if (name == "color") {
+        return output_buffer_color->getHostPointer();
+    }
     throw std::runtime_error("invalid optixgetimg_extra name: " + name);
 }
-
+static void save_exr(float3* ptr, int w, int h, std::string path) {
+    std::vector<float3> data(w * h);
+    std::copy_n(ptr, w * h, data.data());
+    zeno::image_flip_vertical(data.data(), w, h);
+    const char *err = nullptr;
+    int ret = SaveEXR((float *) data.data(), w, h, 3, 1, path.c_str(), &err);
+    if (ret != TINYEXR_SUCCESS) {
+        if (err) {
+            zeno::log_error("failed to perform SaveEXR to {}: {}", path, err);
+            FreeEXRErrorMessage(err);
+        }
+    }
+}
 void optixrender(int fbo, int samples, bool denoise, bool simpleRender) {
     bool imageRendered = false;
     samples = zeno::envconfig::getInt("SAMPLES", samples);
@@ -3439,29 +3467,34 @@ void optixrender(int fbo, int samples, bool denoise, bool simpleRender) {
         auto w = (*output_buffer_o).width();
         auto h = (*output_buffer_o).height();
         stbi_flip_vertically_on_write(true);
-        stbi_write_jpg(path.c_str(), w, h, 4, p, 100);
+        if (zeno::getSession().userData().get2<bool>("output_exr", true)) {
+            auto exr_path = path.substr(0, path.size() - 4) + ".exr";
+            save_exr((float3 *)optixgetimg_extra("color"), w, h, exr_path);
+        }
+        else {
+            stbi_write_jpg(path.c_str(), w, h, 4, p, 100);
+            if (denoise) {
+                const float* _albedo_buffer = reinterpret_cast<float*>(state.albedo_buffer_p.handle);
+                //SaveEXR(_albedo_buffer, w, h, 4, 0, (path+".albedo.exr").c_str(), nullptr);
+                auto a_path = path + ".albedo.pfm";
+                write_pfm(a_path, w, h, _albedo_buffer);
+
+                const float* _normal_buffer = reinterpret_cast<float*>(state.normal_buffer_p.handle);
+                //SaveEXR(_normal_buffer, w, h, 4, 0, (path+".normal.exr").c_str(), nullptr);
+                auto n_path = path + ".normal.pfm";
+                write_pfm(n_path, w, h, _normal_buffer);
+            }
+        }
         zeno::log_info("optix: saving screenshot {}x{} to {}", w, h, path);
         ud.erase("optix_image_path");
 
-        if (denoise) {
-            const float* _albedo_buffer = reinterpret_cast<float*>(state.albedo_buffer_p.handle);
-            //SaveEXR(_albedo_buffer, w, h, 4, 0, (path+".albedo.exr").c_str(), nullptr);
-            auto a_path = path + ".albedo.pfm";
-            write_pfm(a_path, w, h, _albedo_buffer);
-
-            const float* _normal_buffer = reinterpret_cast<float*>(state.normal_buffer_p.handle);
-            //SaveEXR(_normal_buffer, w, h, 4, 0, (path+".normal.exr").c_str(), nullptr);
-            auto n_path = path + ".normal.pfm";
-            write_pfm(n_path, w, h, _normal_buffer); 
-        }
-
         // AOV
         if (zeno::getSession().userData().get2<bool>("output_aov", true)) {
             path = path.substr(0, path.size() - 4);
-            stbi_write_png((path + ".diffuse.png").c_str(), w, h, 4 , optixgetimg_extra("diffuse"), 0);
-            stbi_write_png((path + ".specular.png").c_str(), w, h, 4 , optixgetimg_extra("specular"), 0);
-            stbi_write_png((path + ".transmit.png").c_str(), w, h, 4 , optixgetimg_extra("transmit"), 0);
-            stbi_write_png((path + ".background.png").c_str(), w, h, 4 , optixgetimg_extra("background"), 0);
+            save_exr((float3 *)optixgetimg_extra("diffuse"), w, h, path + ".diffuse.exr");
+            save_exr((float3 *)optixgetimg_extra("specular"), w, h, path + ".specular.exr");
+            save_exr((float3 *)optixgetimg_extra("transmit"), w, h, path + ".transmit.exr");
+            save_exr((float3 *)optixgetimg_extra("background"), w, h, path + ".background.exr");
         }
         imageRendered = true;
     }

zenovis/xinxinoptix/optixPathTracer.h

@@ -57,10 +57,11 @@ struct Params
     float4*      accum_buffer_T;
     float4*      accum_buffer_B;
     uchar4*      frame_buffer;
-    uchar4*      frame_buffer_D;
-    uchar4*      frame_buffer_S;
-    uchar4*      frame_buffer_T;
-    uchar4*      frame_buffer_B;
+    float3*      frame_buffer_C;
+    float3*      frame_buffer_D;
+    float3*      frame_buffer_S;
+    float3*      frame_buffer_T;
+    float3*      frame_buffer_B;
 
     float3*      albedo_buffer;
     float3*      normal_buffer;