Merge pull request #3061 from LouRohanNV:fixUsdMaterialInterfaces

PiperOrigin-RevId: 868177637
Change-Id: I2de8f2af3b5fc126d103256b1ea00c5f8c7a8955

Author: copybara-github

plugin/usd_decoder/material_parsing.cc

@@ -15,6 +15,7 @@
 #include "material_parsing.h"
 
 #include <cstdio>
+#include <filesystem>
 #include <optional>
 #include <string>
 
@@ -24,15 +25,13 @@
 #include <pxr/base/tf/staticData.h>
 #include <pxr/base/tf/staticTokens.h>
 #include <pxr/base/tf/token.h>
-#include <pxr/usd/ar/asset.h>
-#include <pxr/usd/ar/resolvedPath.h>
-#include <pxr/usd/ar/resolver.h>
 #include <pxr/usd/usd/prim.h>
 #include <pxr/usd/usdShade/input.h>
 #include <pxr/usd/usdShade/material.h>
 #include <pxr/usd/usdShade/shader.h>
 #include <pxr/usd/usdShade/types.h>
 #include <pxr/usd/usdShade/udimUtils.h>
+#include <pxr/usd/usdShade/utils.h>
 
 // Using to satisfy TF_DEFINE_PRIVATE_TOKENS macro below and avoid operating in
 // PXR_NS.
@@ -119,84 +118,104 @@ ResolvedShaderInput<T> ReadUsdUVTexture(mjSpec* spec,
     texture->colorspace = mjtColorSpace::mjCOLORSPACE_AUTO;
   }
 
-  pxr::ArResolver& resolver = pxr::ArGetResolver();
   pxr::SdfAssetPath resolved_texture_asset_path;
+
   if (auto file_input = shader.GetInput(kTokens->file)) {
-    file_input.Get(&resolved_texture_asset_path);
+    // Use GetValueProducingAttributes to follow all connections including
+    // material interfaces.
+    auto value_attrs =
+        pxr::UsdShadeUtils::GetValueProducingAttributes(file_input);
+    for (const auto& attr : value_attrs) {
+      if (attr.Get(&resolved_texture_asset_path) &&
+          !resolved_texture_asset_path.GetAssetPath().empty()) {
+        break;
+      }
+    }
   } else {
     mju_error("UsdUVTexture missing inputs:file.");
     return out;
   }
 
-  std::string resolved_texture_path =
-      resolved_texture_asset_path.GetResolvedPath();
+  // Use the resolved path from USD, fall back to asset path if not resolved
+  std::string resolved_path_str = resolved_texture_asset_path.GetResolvedPath();
+  if (resolved_path_str.empty()) {
+    resolved_path_str = resolved_texture_asset_path.GetAssetPath();
+  }
 
-  if (pxr::UsdShadeUdimUtils::IsUdimIdentifier(resolved_texture_path)) {
-    mju_error("MuJoCo does not support UDIM textures: %s",
-              resolved_texture_path.c_str());
+  if (resolved_path_str.empty()) {
+    mju_warning("UsdUVTexture %s: No texture file path specified.",
+                shader.GetPath().GetAsString().c_str());
     return out;
   }
 
-  auto extension = resolver.GetExtension(resolved_texture_path);
+  if (pxr::UsdShadeUdimUtils::IsUdimIdentifier(resolved_path_str)) {
+    mju_error("MuJoCo does not support UDIM textures: %s",
+              resolved_path_str.c_str());
+    return out;
+  }
 
-  FILE* fp = fopen(resolved_texture_path.c_str(), "r");
-  if (fp == nullptr) {
-    mju_error(
+  if (!std::filesystem::exists(resolved_path_str)) {
+    mju_warning(
         "USD decoder only supports assets that are available on the file "
-        "system");
+        "system. Could not open: %s",
+        resolved_path_str.c_str());
     return out;
   }
+
   texture->nchannel = nchannels;
-  mjs_setString(texture->file, resolved_texture_path.c_str());
+  mjs_setString(texture->file, resolved_path_str.c_str());
   out.sampler = texture;
   return out;
 }
 
-// Given an input to a shader, attempts to bake the shader and it's inputs
-// into a singular value or texture sampler (ResolvedShaderInput).
-template <typename T>
-ResolvedShaderInput<T> ReadShaderInput(
-    mjSpec* spec, const pxr::UsdShadeConnectableAPI source,
-    const pxr::TfToken source_name,
-    const pxr::UsdShadeAttributeType source_type) {
-  ResolvedShaderInput<T> out;
-  pxr::UsdShadeShader shader(source.GetPrim());
-  pxr::TfToken shader_id = GetShaderId(shader);
-  if (shader_id == kTokens->UsdUVTexture) {
-    unsigned nchannels = -1;
-    if (source_name == kTokens->rgb) {
-      nchannels = 3;
-    } else if (source_name == kTokens->r) {
-      nchannels = 1;
-    } else {
-      mju_error("Unsupported texture channel: %s", source_name.GetText());
-      return out;
-    }
-
-    return ReadUsdUVTexture<T>(spec, shader, nchannels);
-  } else {
-    mju_warning("Unsupported shader type: %s", shader_id.GetText());
-  }
-
-  return out;
-}
-
 // Reads UsdShadeInput as a value of type T or a texture sampler.
+// Uses GetValueProducingAttributes to follow all connections including
+// material interfaces.
 template <typename T>
 ResolvedShaderInput<T> ReadShaderInput(mjSpec* spec, pxr::UsdShadeInput input) {
   ResolvedShaderInput<T> out;
   if (!input.GetPrim().IsValid()) {
     return out;
   }
 
-  pxr::UsdShadeConnectableAPI source;
-  pxr::TfToken source_name;
-  pxr::UsdShadeAttributeType source_type;
-  if (input.GetConnectedSource(&source, &source_name, &source_type)) {
-    out = ReadShaderInput<T>(spec, source, source_name, source_type);
-  } else {
-    out.value = ReadInput<T>(input);
+  // GetValueProducingAttributes follows connections recursively, including
+  // through material interfaces, and returns the attributes that produce
+  // values.
+  auto value_attrs = pxr::UsdShadeUtils::GetValueProducingAttributes(input);
+
+  for (const auto& attr : value_attrs) {
+    // Check if this attribute belongs to a UsdUVTexture shader (for textures)
+    pxr::UsdPrim prim = attr.GetPrim();
+    pxr::UsdShadeShader shader(prim);
+    if (shader) {
+      pxr::TfToken shader_id = GetShaderId(shader);
+      if (shader_id == kTokens->UsdUVTexture) {
+        // Determine channel count from the output name
+        pxr::TfToken attr_name = attr.GetBaseName();
+        unsigned nchannels = 3;  // default
+        if (attr_name == kTokens->rgb) {
+          nchannels = 3;
+        } else if (attr_name == kTokens->r) {
+          nchannels = 1;
+        }
+        out = ReadUsdUVTexture<T>(spec, shader, nchannels);
+        if (out.sampler.has_value()) {
+          return out;
+        }
+      }
+      // For other shader types, fall through to try reading as a value
+    }
+
+    // Try to read as a direct value
+    out.value = ReadInput<T>(pxr::UsdShadeInput(attr));
+    if (out.value.has_value()) {
+      return out;
+    }
   }
+
+  // If no value-producing attributes found, try reading the input directly
+  // (for non-connected inputs with authored values)
+  out.value = ReadInput<T>(input);
   return out;
 }