Return cryptomatte metadatas in GetRenderStats #1164 (#1178)

Autodesk · Jun 15, 2022 · e04a694 · e04a694
1 parent 5f1993a
commit e04a694
Show file tree

Hide file tree

Showing 4 changed files with 71 additions and 0 deletions.
diff --git a/common/constant_strings.h b/common/constant_strings.h
@@ -94,6 +94,10 @@ ASTR(ArnoldOptions);
 ASTR(BOOL);
 ASTR(BYTE);
 ASTR(CPU);
+ASTR(crypto_asset);
+ASTR(crypto_material);
+ASTR(crypto_object);
+ASTR(custom_attributes);
 ASTR(FLOAT);
 ASTR(GI_diffuse_depth);
 ASTR(GI_diffuse_samples);

diff --git a/render_delegate/render_delegate.cpp b/render_delegate/render_delegate.cpp
@@ -848,6 +848,42 @@ VtDictionary HdArnoldRenderDelegate::GetRenderStats() const
     AiRenderGetHintFlt(str::total_progress, total_progress);
 #endif
     stats[_tokens->percentDone] = total_progress;
+
+    // If there are cryptomatte drivers, we look for the metadata that is stored in each of them.
+    // In theory, we could just look for the first driver, but for safety we're doing it for all of them
+    for (const auto& cryptoDriver : _cryptomatteDrivers) {
+        const AtNode *driver = AiNodeLookUpByName(_universe, cryptoDriver);
+        if (!driver)
+            continue;
+        if (AiNodeLookUpUserParameter(driver, str::custom_attributes) == nullptr)
+            continue;
+        const AtArray *customAttrsArray = AiNodeGetArray(driver, str::custom_attributes);
+        if (customAttrsArray == nullptr)
+            continue;
+        unsigned int customAttrsCount = AiArrayGetNumElements(customAttrsArray);
+        for (unsigned int i = 0; i < customAttrsCount; ++i) {
+            AtString customAttr = AiArrayGetStr(customAttrsArray, i);
+            std::string customAttrStr(customAttr.c_str());
+            // the custom_attributes attribute will be an array of strings, where each  
+            // element is set like:
+            // "STRING cryptomatte/f834d0a/conversion uint32_to_float32"
+            // where the second element is the metadata name and the last one
+            // is the metadata value
+            size_t pos = customAttrStr.find_first_of(' ');
+            if (pos == std::string::npos)
+                continue;
+            std::string customAttrType = customAttrStr.substr(0, pos);
+            customAttrStr = customAttrStr.substr(pos + 1);
+
+            pos = customAttrStr.find_first_of(' ');
+            if (pos == std::string::npos)
+                continue;
+            std::string metadataName = customAttrStr.substr(0, pos);
+            std::string metadataVal = customAttrStr.substr(pos + 1);
+            // TODO do we want to check if the metadata is not a string ?
+            stats[TfToken(metadataName)] = TfToken(metadataVal);
+        }
+    }
     return stats;
 }
 
@@ -1427,5 +1463,15 @@ AtNode* HdArnoldRenderDelegate::GetSubdivDicingCamera(HdRenderIndex* renderIndex
     return AiNodeLookUpByName(_universe, AtString(_subdiv_dicing_camera.GetText()));
 }
 
+void HdArnoldRenderDelegate::RegisterCryptomatteDriver(const AtString& driver)
+{
+    _cryptomatteDrivers.insert(driver);
+
+}
+void HdArnoldRenderDelegate::ClearCryptomatteDrivers()
+{
+   _cryptomatteDrivers.clear();
+}
+
 
 PXR_NAMESPACE_CLOSE_SCOPE
diff --git a/render_delegate/render_delegate.h b/render_delegate/render_delegate.h
@@ -462,6 +462,16 @@ class HdArnoldRenderDelegate final : public HdRenderDelegate {
     HDARNOLD_API
     std::vector<AtNode*> GetAovShaders(HdRenderIndex* renderIndex);
 
+    // Store the list of cryptomatte driver names, so that we can get the cryptomatte
+    // metadatas in their attribute "custom_attributes"
+    /// @param driver Name of a driver used for a cryptomatte AOVs (crypto_material, crypto_asset, crypto_object)
+    HDARNOLD_API
+    void RegisterCryptomatteDriver(const AtString& driver);
+
+    // Clear the list of cryptomatte driver names, before outputs are setup
+    HDARNOLD_API
+    void ClearCryptomatteDrivers();
+
     /// Get the current Window NDC, as a resolution-independant value, 
     /// defaulting to (0,0,1,1)
     ///
@@ -563,6 +573,7 @@ class HdArnoldRenderDelegate final : public HdRenderDelegate {
     int _verbosityLogFlags = AI_LOG_WARNINGS | AI_LOG_ERRORS;
     bool _ignoreVerbosityLogFlags = false;
     bool _isArnoldActive = false;
+    std::unordered_set<AtString, AtStringHash> _cryptomatteDrivers;
 };
 
 PXR_NAMESPACE_CLOSE_SCOPE
diff --git a/render_delegate/render_pass.cpp b/render_delegate/render_pass.cpp
@@ -706,6 +706,7 @@ void HdArnoldRenderPass::_Execute(const HdRenderPassStateSharedPtr& renderPassSt
             _usingFallbackBuffers = false;
             renderParam->Interrupt();
             _ClearRenderBuffers();
+            _renderDelegate->ClearCryptomatteDrivers();
             AiNodeSetPtr(_mainDriver, str::color_pointer, nullptr);
             AiNodeSetPtr(_mainDriver, str::depth_pointer, nullptr);
             AiNodeSetPtr(_mainDriver, str::id_pointer, nullptr);
@@ -799,6 +800,15 @@ void HdArnoldRenderPass::_Execute(const HdRenderPassStateSharedPtr& renderPassSt
                     } else {
                         aovName = sourceName.c_str();
                     }
+                    // If this driver is meant for one of the cryptomatte AOVs, it will be filled with the 
+                    // cryptomatte metadatas through the user data "custom_attributes". We want to store 
+                    // the driver node names in the render delegate, so that we can lookup this user data
+                    // during GetRenderStats
+                    if (binding.aovName == str::t_crypto_asset || 
+                        binding.aovName == str::t_crypto_material ||
+                        binding.aovName == str::t_crypto_object)
+                        _renderDelegate->RegisterCryptomatteDriver(driverNameStr);
+
                     output = AtString{
                         TfStringPrintf(
                             "%s %s %s %s", aovName, arnoldTypes.outputString, filterName, AiNodeGetName(buffer.driver))