Save extra AOVs

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

@@ -114,6 +114,7 @@ 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;
@@ -222,7 +223,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 +238,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
@@ -646,6 +649,16 @@ namespace
             }
             else if (extension == ".png")
             {
+                // Test if the image contains appropriate data
+                if (image.properties().m_pixel_format != PixelFormatHalf &&
+                    image.properties().m_channel_count != 4)
+                {
+                    RENDERER_LOG_ERROR(
+                        "failed to write image file %s: data images can only be saved as exr.",
+                        file_path);
+                    return false;
+                }
+
                 write_png_image(
                     bf_file_path,
                     image,
@@ -687,6 +700,38 @@ namespace
 
         return true;
     }
+
+    bool write_extra_aovs(
+        const ImageStack & images,
+        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 = images.size(); i < e; ++i)
+        {
+            const Image & image = images.get_image(i);
+
+            // Compute AOV image file path.
+            const string aov_name = images.get_name(i);
+            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,7 +765,7 @@ bool Frame::write_aov_images(const char* file_path) const
 
     bool success = true;
 
-    if (!aovs().empty())
+    if (!aovs().empty() || !aov_images().empty())
     {
         const bf::path boost_file_path(file_path);
         const bf::path directory = boost_file_path.parent_path();
@@ -741,6 +786,11 @@ bool Frame::write_aov_images(const char* file_path) const
             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(), directory, base_file_name, extension);
+        }
     }
 
     return success;
@@ -776,6 +826,16 @@ bool Frame::write_main_and_aov_images() const
         }
     }
 
+    if (impl->m_save_extra_aovs && !aov_images().empty())
+    {
+        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(), directory, base_file_name, extension);
+    }
+
     return success;
 }
 
@@ -795,12 +855,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 +882,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 +1028,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 +1219,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;
 }