Merge pull request #1389 from johnhaddon/usdSceneStageCache

USDScene : Support reading from in-memory stages via UsdUtilsStageCache
ImageEngine · Sep 15, 2023 · 1afc717 · 1afc717
2 parents 54c5f6c + 8e3f071
commit 1afc717
Show file tree

Hide file tree

Showing 4 changed files with 30 additions and 2 deletions.
diff --git a/Changes b/Changes
@@ -5,6 +5,7 @@ Improvements
 ------------
 
 - USD PrimitiveAlgo : Added `readPrimitiveVariable()` utility method for reading from regular `UsdAttributes`.
+- USDScene : Added support for reading from in-memory stages by passing a filename of the form `stageCache:{id}.usd` where `{id}` specifies a stage which has been inserted in the `UsdUtilsStageCache`.
 
 Fixes
 -----

diff --git a/SConstruct b/SConstruct
@@ -2967,6 +2967,7 @@ else :
 		"usdLux",
 		"usdSkel",
 		"usdShade",
+		"usdUtils",
 		"sdf",
 		"tf",
 		"pcp",

diff --git a/contrib/IECoreUSD/src/IECoreUSD/USDScene.cpp b/contrib/IECoreUSD/src/IECoreUSD/USDScene.cpp
@@ -70,6 +70,7 @@ IECORE_PUSH_DEFAULT_VISIBILITY
 #include "pxr/usd/usdShade/material.h"
 #include "pxr/usd/usdShade/materialBindingAPI.h"
 #include "pxr/usd/usdShade/connectableAPI.h"
+#include "pxr/usd/usdUtils/stageCache.h"
 #ifdef IECOREUSD_WITH_OPENVDB
 #include "pxr/usd/usdVol/fieldBase.h"
 #endif
@@ -84,6 +85,7 @@ IECORE_POP_DEFAULT_VISIBILITY
 
 #include "tbb/concurrent_hash_map.h"
 
+#include <filesystem>
 #include <iostream>
 #include <mutex>
 
@@ -688,10 +690,24 @@ class USDScene::IO : public RefCounted
 			switch( openMode )
 			{
 				case IndexedIO::Read : {
-					pxr::UsdStageRefPtr stage = pxr::UsdStage::Open( fileName );
+					static const std::string g_stageCachePrefix( "stageCache:" );
+					pxr::UsdStageRefPtr stage;
+					if( boost::starts_with( fileName, g_stageCachePrefix ) )
+					{
+						// Get Id from filename of form "stageCache:{id}.usd"
+						std::filesystem::path path( fileName.substr( g_stageCachePrefix.size() ) );
+						path.replace_extension();
+						stage = pxr::UsdUtilsStageCache::Get().Find(
+							pxr::UsdStageCache::Id::FromString( path.string() )
+						);
+					}
+					else
+					{
+						stage = pxr::UsdStage::Open( fileName );
+					}
 					if( !stage )
 					{
-						throw IECore::Exception( boost::str( boost::format( "USDScene : Failed to open USD file: '%1%'" ) % fileName ) );
+						throw IECore::Exception( boost::str( boost::format( "USDScene : Failed to open USD stage : '%1%'" ) % fileName ) );
 					}
 					return stage;
 				}

diff --git a/contrib/IECoreUSD/test/IECoreUSD/USDSceneTest.py b/contrib/IECoreUSD/test/IECoreUSD/USDSceneTest.py
@@ -3911,5 +3911,15 @@ def testUnconnectedMaterialOutput( self ) :
 		self.assertNotIn( "cycles:surface", sphere.attributeNames() )
 		self.assertIsNone( sphere.readAttribute( "cycles:surface", 0 ) )
 
+	def testReadFromStageCache( self ) :
+
+		stage = pxr.Usd.Stage.CreateInMemory()
+		pxr.UsdGeom.Sphere.Define( stage, "/sphere" )
+		id = pxr.UsdUtils.StageCache.Get().Insert( stage )
+
+		root = IECoreScene.SceneInterface.create( "stageCache:{}.usd".format( id.ToString() ), IECore.IndexedIO.OpenMode.Read )
+		self.assertEqual( root.childNames(), [ "sphere" ] )
+		self.assertIsInstance( root.child( "sphere" ).readObject( 0 ), IECoreScene.SpherePrimitive )
+
 if __name__ == "__main__":
 	unittest.main()