Save extra AOVs

src/appleseed/renderer/modeling/frame/frame.cpp

@@ -114,13 +114,15 @@ struct Frame::Impl
     bool                    m_render_stamp_enabled;
     string                  m_render_stamp_format;
     DenoisingMode           m_denoising_mode;
+    bool                    m_save_extra_aovs;
 
     // Images.
     unique_ptr<Image>       m_image;
     unique_ptr<ImageStack>  m_aov_images;
     AOVContainer            m_aovs;
     AOVContainer            m_internal_aovs;
     DenoiserAOV*            m_denoiser_aov;
+	vector<size_t>          m_extra_aovs;
 };
 
 Frame::Frame(
@@ -222,7 +224,8 @@ void Frame::print_settings()
         "  filter size                   %f\n"
         "  crop window                   (%s, %s)-(%s, %s)\n"
         "  denoising mode                %s\n"
-        "  render stamp                  %s",
+        "  render stamp                  %s\n"
+        "  save extra aovs               %s",
         camera_name ? camera_name : "none",
         pretty_uint(impl->m_frame_width).c_str(),
         pretty_uint(impl->m_frame_height).c_str(),
@@ -236,7 +239,8 @@ void Frame::print_settings()
         pretty_uint(impl->m_crop_window.max[1]).c_str(),
         impl->m_denoising_mode == DenoisingMode::Off ? "off" :
         impl->m_denoising_mode == DenoisingMode::WriteOutputs ? "write outputs" : "denoise",
-        impl->m_render_stamp_enabled ? "enabled" : "disabled");
+        impl->m_render_stamp_enabled ? "enabled" : "disabled",
+        impl->m_save_extra_aovs ? "enabled" : "disabled");
 }
 
 const char* Frame::get_active_camera_name() const
@@ -280,9 +284,12 @@ const AOVContainer& Frame::internal_aovs() const
 
 size_t Frame::create_extra_aov_image(const char* name) const
 {
-    const size_t index = aov_images().get_index(name);
+    size_t index = aov_images().get_index(name);
     if (index == size_t(~0) && aov_images().size() < MaxAOVCount)
-        return aov_images().append(name, 4, PixelFormatFloat);
+    {
+        index = aov_images().append(name, 4, PixelFormatFloat);
+        impl->m_extra_aovs.push_back(index);
+    }
 
     return index;
 }
@@ -687,6 +694,42 @@ namespace
 
         return true;
     }
+
+    bool write_extra_aovs(
+        const ImageStack&       images,
+            const vector<size_t>&   aov_indices,
+        const bf::path&         directory,
+        const string&           base_file_name,
+        const string&           extension)
+    {
+        bool success = true;
+
+        if (extension != ".exr")
+        {
+            RENDERER_LOG_ERROR("extra AOVs can only be saved as exr.");
+            return false;
+        }
+
+        for (size_t i = 0, e = aov_indices.size(); i < e; ++i)
+        {
+            size_t image_index = aov_indices[i];
+            assert(image_index < images.size());
+
+            const Image & image = images.get_image(image_index);
+
+            // Compute image file path.
+            const string aov_name = images.get_name(image_index);
+            const string safe_aov_name = make_safe_filename(aov_name);
+            const string aov_file_name = base_file_name + "." + safe_aov_name + extension;
+            const string aov_file_path = (directory / aov_file_name).string();
+
+            // Write AOV image.
+            if (!write_image(aov_file_path.c_str(), image))
+                success = false;
+        }
+
+        return success;
+    }
 }
 
 bool Frame::write_main_image(const char* file_path) const
@@ -720,27 +763,34 @@ bool Frame::write_aov_images(const char* file_path) const
 
     bool success = true;
 
-    if (!aovs().empty())
+    const bf::path boost_file_path(file_path);
+    const bf::path directory = boost_file_path.parent_path();
+    const string base_file_name = boost_file_path.stem().string();
+    const string extension = boost_file_path.extension().string();
+
+    for (size_t i = 0, e = aovs().size(); i < e; ++i)
     {
-        const bf::path boost_file_path(file_path);
-        const bf::path directory = boost_file_path.parent_path();
-        const string base_file_name = boost_file_path.stem().string();
-        const string extension = boost_file_path.extension().string();
+        const AOV* aov = aovs().get_by_index(i);
 
-        for (size_t i = 0, e = aovs().size(); i < e; ++i)
-        {
-            const AOV* aov = aovs().get_by_index(i);
+        // Compute AOV image file path.
+        const string aov_name = aov->get_name();
+        const string safe_aov_name = make_safe_filename(aov_name);
+        const string aov_file_name = base_file_name + "." + safe_aov_name + extension;
+        const string aov_file_path = (directory / aov_file_name).string();
 
-            // Compute AOV image file path.
-            const string aov_name = aov->get_name();
-            const string safe_aov_name = make_safe_filename(aov_name);
-            const string aov_file_name = base_file_name + "." + safe_aov_name + extension;
-            const string aov_file_path = (directory / aov_file_name).string();
+        // Write AOV image.
+        if (!write_image(aov_file_path.c_str(), aov->get_image(), aov))
+            success = false;
+    }
 
-            // Write AOV image.
-            if (!write_image(aov_file_path.c_str(), aov->get_image(), aov))
-                success = false;
-        }
+    if (impl->m_save_extra_aovs && !aov_images().empty())
+    {
+        success = success && write_extra_aovs(
+            aov_images(),
+            impl->m_extra_aovs,
+            directory,
+            base_file_name,
+            extension);
     }
 
     return success;
@@ -776,6 +826,21 @@ bool Frame::write_main_and_aov_images() const
         }
     }
 
+    if (impl->m_save_extra_aovs)
+    {
+        const bf::path boost_file_path(filepath);
+        const bf::path directory = boost_file_path.parent_path();
+        const string base_file_name = boost_file_path.stem().string();
+        const string extension = boost_file_path.extension().string();
+
+        success = success && write_extra_aovs(
+            aov_images(),
+            impl->m_extra_aovs,
+            directory,
+            base_file_name,
+            extension);
+    }
+
     return success;
 }
 
@@ -795,12 +860,13 @@ void Frame::write_main_and_aov_images_to_multipart_exr(const char* file_path) co
 
     writer.begin_multipart_exr();
 
+    static const char* ChannelNames[] = { "R", "G", "B", "A" };
+
     // Always save the main image as half floats.
     {
         const Image& image = *impl->m_image;
         const CanvasProperties& props = image.properties();
         images.emplace_back(image, props.m_tile_width, props.m_tile_height, PixelFormatHalf);
-        static const char* ChannelNames[] = { "R", "G", "B", "A" };
         writer.append_part("beauty", images.back(), image_attributes, 4, ChannelNames);
     }
 
@@ -821,6 +887,18 @@ void Frame::write_main_and_aov_images_to_multipart_exr(const char* file_path) co
             writer.append_part(aov_name.c_str(), image, image_attributes, aov->get_channel_count(), aov->get_channel_names());
     }
 
+    if (impl->m_save_extra_aovs)
+    {
+        for (size_t i = 0, e = aov_images().size(); i < e; ++i)
+        {
+            const Image & image = aov_images().get_image(i);
+            const CanvasProperties& props = image.properties();
+            const string aov_name = aov_images().get_name(i);
+            assert(props.m_channel_count == 4);
+            writer.append_part(aov_name.c_str(), image, image_attributes, props.m_channel_count, ChannelNames);
+        }
+    }
+
     create_parent_directories(file_path);
     writer.write_multipart_exr(file_path);
 
@@ -955,6 +1033,9 @@ void Frame::extract_parameters()
             impl->m_denoising_mode = DenoisingMode::Off;
         }
     }
+
+    // Retrieve save extra AOVs parameter
+    impl->m_save_extra_aovs = m_params.get_optional<bool>("save_extra_aovs", false);
 }
 
 
@@ -1143,6 +1224,14 @@ DictionaryArray FrameFactory::get_input_metadata()
                 Dictionary()
                     .insert("enable_render_stamp", "true")));
 
+    metadata.push_back(
+        Dictionary()
+            .insert("name", "save_extra_aovs")
+            .insert("label", "Save Extra AOVs")
+            .insert("type", "boolean")
+            .insert("use", "optional")
+            .insert("default", "false"));
+
     return metadata;
 }