USD: Add new get_prim_map API callable from python on_import hooks

When importing an USD, the Blender object names do not necessarily match
Prim names:

```
#usda 1.0
def Xform "xform1"
{
  def Mesh "MyObject"  # Will be imported as "MyObject"
}
def Xform "xform2"
{
  def Mesh "MyObject"  # Will be imported as "MyObject.001"
}
def Xform "xform2"
{
  def Mesh "MyObject"  # Will be imported as "MyObject.002"
}
```

This makes it difficult in the [USD Import Hooks]
(https://docs.blender.org/api/current/bpy.types.USDHook.html) to link a
Blender object back to its source Prim for additional processing. A
typical use cases for games is to generate UVs, create and apply a
material on the fly when importing a collision shape that does not have
a visual representation (hence no materials) based on some Prim
attributes, but that the artist needs to differenciate in Blender's
viewport.

The Import context exposes a new method `get_prim_map()` that returns a
`dict` of `prim path` / `list of Blender ID`.

For example, given the following USD scene,
```
/
 |--XformThenCube [def Xform]
 |   `--Cube [def Cube]
 |--XformThenXformCube [def Xform]
 |   `--XformIntermediate [def Xform]
 |       `--Cube [def Mesh]
 |--Cube [def Cube]
 `--Material [def Material]
     `--Principled_BSDF [def Shader]
```

the `get_prim_map()` method will return a map as:

```python
@static_method
def on_import(import_context):
  pprint(import_context.get_prim_map())
```

```json
{
  "/Cube": [bpy.data.objects["Cube.002"], bpy.data.meshes["Cube.002"]],
  "/XformThenCube": [bpy.data.objects["XformThenCube"]],
  "/XformThenCube/Cube": [bpy.data.objects["Cube"], bpy.data.meshes["Cube"]],
  "/XformThenXformCube": [bpy.data.objects["XformThenXformCube"]],
  "/XformThenXformCube/XformIntermediate": [bpy.data.objects["XformIntermediate"]],
  "/XformThenXformCube/XformIntermediate/Cube": [bpy.data.objects["Cube.001"], bpy.data.meshes["Cube.001"]],
  "/Material": [bpy.data.materials["Material"]],
}
```

Co-authored-by: Odréanne Breton <odreanne.breton@ubisoft.com>
Co-authored-by: Sttevan Carnali Joga <sttevan.carnali-joga@ubisoft.com>
Co-authored-by: Charles Flèche <charles.fleche@ubisoft.com>

doc/python_api/examples/bpy.types.USDHook.py

@@ -35,6 +35,8 @@ Hook function ``on_import()`` is called after the USD import finalizes. This fun
 as an argument an instance of an internally defined class ``USDSceneImportContext`` which provides the
 following accessors to the scene data:
 
+- ``get_prim_map()`` returns a dict where the key is an imported USD Prim path and the value a list of
+   the IDs created by the imported prim.
 - ``get_stage()`` returns the USD stage which was imported.
 
 The hook functions should return ``True`` on success or ``False`` if the operation was bypassed or

source/blender/io/usd/intern/usd_capi_import.cc

@@ -36,6 +36,7 @@
 #include "DNA_collection_types.h"
 #include "DNA_layer_types.h"
 #include "DNA_listBase.h"
+#include "DNA_material_types.h"
 #include "DNA_object_types.h"
 #include "DNA_scene_types.h"
 #include "DNA_windowmanager_types.h"
@@ -44,6 +45,8 @@
 
 #include "MEM_guardedalloc.h"
 
+#include "RNA_access.hh"
+
 #include "WM_api.hh"
 #include "WM_types.hh"
 
@@ -169,6 +172,7 @@ struct ImportJobData {
   ImportSettings settings;
 
   USDStageReader *archive;
+  ImportedPrimMap prim_map;
 
   bool *stop;
   bool *do_update;
@@ -345,9 +349,19 @@ static void import_startjob(void *customdata, wmJobWorkerStatus *worker_status)
     }
 
     Object *ob = reader->object();
+    if (!ob) {
+      continue;
+    }
 
     reader->read_object_data(data->bmain, 0.0);
 
+    data->prim_map[reader->object_prim_path()].push_back(RNA_id_pointer_create(&ob->id));
+
+    if (ob->data) {
+      data->prim_map[reader->data_prim_path()].push_back(
+          RNA_id_pointer_create(static_cast<ID *>(ob->data)));
+    }
+
     USDPrimReader *parent = reader->parent();
 
     if (parent == nullptr) {
@@ -366,6 +380,14 @@ static void import_startjob(void *customdata, wmJobWorkerStatus *worker_status)
     }
   }
 
+  data->settings.usd_path_to_mat_name.foreach_item(
+      [&](const std::string &path, const std::string &name) {
+        Material *mat = data->settings.mat_name_to_mat.lookup_default(name, nullptr);
+        if (mat) {
+          data->prim_map[path].push_back(RNA_id_pointer_create(&mat->id));
+        }
+      });
+
   if (data->params.import_skeletons) {
     archive->process_armature_modifiers();
   }
@@ -457,7 +479,7 @@ static void import_endjob(void *customdata)
     /* Ensure Python types for invoking hooks are registered. */
     register_hook_converters();
 
-    call_import_hooks(data->archive->stage(), data->params.worker_status->reports);
+    call_import_hooks(data->archive->stage(), data->prim_map, data->params.worker_status->reports);
 
     if (data->is_background_job) {
       /* Blender already returned from the import operator, so we need to store our own extra undo

source/blender/io/usd/intern/usd_hook.cc

@@ -7,6 +7,7 @@
 
 #include "BLI_utildefines.h"
 
+#include "BKE_idtype.hh"
 #include "BKE_report.hh"
 
 #include "DNA_windowmanager_types.h"
@@ -122,14 +123,47 @@ struct USDSceneExportContext {
 struct USDSceneImportContext {
   USDSceneImportContext() = default;
 
-  USDSceneImportContext(pxr::UsdStageRefPtr in_stage) : stage(in_stage) {}
+  USDSceneImportContext(pxr::UsdStageRefPtr in_stage, const ImportedPrimMap &in_prim_map)
+      : stage(in_stage), prim_map(in_prim_map)
+  {
+  }
+
+  void release()
+  {
+    if (prim_map_dict) {
+      delete prim_map_dict;
+    }
+  }
 
   pxr::UsdStageRefPtr get_stage() const
   {
     return stage;
   }
 
+  boost::python::dict get_prim_map()
+  {
+    if (!prim_map_dict) {
+      prim_map_dict = new boost::python::dict;
+
+      for (auto &[path, ids] : prim_map) {
+        if (!prim_map_dict->has_key(path)) {
+          (*prim_map_dict)[path] = boost::python::list();
+        }
+        boost::python::list list = boost::python::extract<boost::python::list>(
+            (*prim_map_dict)[path]);
+
+        for (auto &ptr_rna : ids) {
+          list.append(ptr_rna);
+        }
+      }
+    }
+
+    return *prim_map_dict;
+  }
+
   pxr::UsdStageRefPtr stage;
+  ImportedPrimMap prim_map;
+  boost::python::dict *prim_map_dict = nullptr;
 };
 
 /* Encapsulate arguments for material export. */
@@ -181,7 +215,8 @@ void register_hook_converters()
       .def("get_stage", &USDMaterialExportContext::get_stage);
 
   python::class_<USDSceneImportContext>("USDSceneImportContext")
-      .def("get_stage", &USDSceneImportContext::get_stage);
+      .def("get_stage", &USDSceneImportContext::get_stage)
+      .def("get_prim_map", &USDSceneImportContext::get_prim_map);
 
   PyGILState_Release(gilstate);
 }
@@ -209,13 +244,14 @@ static void handle_python_error(USDHook *hook, ReportList *reports)
 class USDHookInvoker {
  public:
   /* Attempt to call the function, if defined by the registered hooks. */
-  void call() const
+  void call()
   {
     if (hook_list().empty()) {
       return;
     }
 
     PyGILState_STATE gilstate = PyGILState_Ensure();
+    init_in_gil();
 
     /* Iterate over the hooks and invoke the hook function, if it's defined. */
     USDHookList::const_iterator hook_iter = hook_list().begin();
@@ -251,6 +287,7 @@ class USDHookInvoker {
       }
     }
 
+    release_in_gil();
     PyGILState_Release(gilstate);
   }
 
@@ -263,6 +300,9 @@ class USDHookInvoker {
    * python::call_method<void>(hook_obj, function_name(), arg1, arg2); */
   virtual void call_hook(PyObject *hook_obj) const = 0;
 
+  virtual void init_in_gil(){};
+  virtual void release_in_gil(){};
+
   /* Reports list provided when constructing the subclass, used by #call() to store reports. */
   ReportList *reports_;
 };
@@ -325,7 +365,8 @@ class OnImportInvoker : public USDHookInvoker {
   USDSceneImportContext hook_context_;
 
  public:
-  OnImportInvoker(pxr::UsdStageRefPtr stage, ReportList *reports) : hook_context_(stage)
+  OnImportInvoker(pxr::UsdStageRefPtr stage, const ImportedPrimMap &prim_map, ReportList *reports)
+      : hook_context_(stage, prim_map)
   {
     reports_ = reports;
   }
@@ -340,6 +381,11 @@ class OnImportInvoker : public USDHookInvoker {
   {
     python::call_method<bool>(hook_obj, function_name(), REF(hook_context_));
   }
+
+  void release_in_gil() override
+  {
+    hook_context_.release();
+  }
 };
 
 void call_export_hooks(pxr::UsdStageRefPtr stage, Depsgraph *depsgraph, ReportList *reports)
@@ -365,13 +411,15 @@ void call_material_export_hooks(pxr::UsdStageRefPtr stage,
   on_material_export.call();
 }
 
-void call_import_hooks(pxr::UsdStageRefPtr stage, ReportList *reports)
+void call_import_hooks(pxr::UsdStageRefPtr stage,
+                       const ImportedPrimMap &prim_map,
+                       ReportList *reports)
 {
   if (hook_list().empty()) {
     return;
   }
 
-  OnImportInvoker on_import(stage, reports);
+  OnImportInvoker on_import(stage, prim_map, reports);
   on_import.call();
 }
 

source/blender/io/usd/intern/usd_hook.hh

@@ -8,10 +8,13 @@
 
 struct Depsgraph;
 struct Material;
+struct PointerRNA;
 struct ReportList;
 
 namespace blender::io::usd {
 
+using ImportedPrimMap = std::map<std::string, std::vector<PointerRNA>>;
+
 /** Ensure classes and type converters necessary for invoking import and export hooks
  * are registered. */
 void register_hook_converters();
@@ -26,6 +29,8 @@ void call_material_export_hooks(pxr::UsdStageRefPtr stage,
                                 ReportList *reports);
 
 /** Call the 'on_import' chaser function defined in the registered USDHook classes. */
-void call_import_hooks(pxr::UsdStageRefPtr stage, ReportList *reports);
+void call_import_hooks(pxr::UsdStageRefPtr stage,
+                       const ImportedPrimMap &imported_id_links,
+                       ReportList *reports);
 
 }  // namespace blender::io::usd

source/blender/io/usd/intern/usd_reader_prim.hh

@@ -121,6 +121,16 @@ class USDPrimReader {
     return prim_path_;
   }
 
+  virtual std::string object_prim_path() const
+  {
+    return prim_path();
+  }
+
+  virtual std::string data_prim_path() const
+  {
+    return prim_path();
+  }
+
   void set_is_in_instancer_proto(bool flag)
   {
     is_in_instancer_proto_ = flag;

source/blender/io/usd/intern/usd_reader_stage.hh

@@ -38,7 +38,7 @@ class USDStageReader {
  protected:
   pxr::UsdStageRefPtr stage_;
   USDImportParams params_;
-  ImportSettings settings_;
+  const ImportSettings &settings_;
 
   blender::Vector<USDPrimReader *> readers_;
 

source/blender/io/usd/intern/usd_reader_xform.cc

@@ -59,6 +59,11 @@ void USDXformReader::read_object_data(Main * /*bmain*/, const double motionSampl
   set_props(use_parent_xform(), motionSampleTime);
 }
 
+std::string USDXformReader::object_prim_path() const
+{
+  return get_xformable().GetPrim().GetPath().GetAsString();
+}
+
 void USDXformReader::read_matrix(float r_mat[4][4] /* local matrix */,
                                  const float time,
                                  const float scale,
@@ -150,8 +155,7 @@ bool USDXformReader::is_root_xform_prim() const
 
 std::optional<XformResult> USDXformReader::get_local_usd_xform(const float time) const
 {
-  pxr::UsdGeomXformable xformable = use_parent_xform_ ? pxr::UsdGeomXformable(prim_.GetParent()) :
-                                                        pxr::UsdGeomXformable(prim_);
+  pxr::UsdGeomXformable xformable = get_xformable();
 
   if (!xformable) {
     /* This might happen if the prim is a Scope. */
@@ -172,4 +176,9 @@ std::optional<XformResult> USDXformReader::get_local_usd_xform(const float time)
   return XformResult(pxr::GfMatrix4f(xform), is_constant);
 }
 
+pxr::UsdGeomXformable USDXformReader::get_xformable() const
+{
+  pxr::UsdPrim prim = use_parent_xform_ ? prim_.GetParent() : prim_;
+  return pxr::UsdGeomXformable(prim);
+}
 }  // namespace blender::io::usd

source/blender/io/usd/intern/usd_reader_xform.hh

@@ -39,6 +39,8 @@ class USDXformReader : public USDPrimReader {
   void create_object(Main *bmain, double motionSampleTime) override;
   void read_object_data(Main *bmain, double motionSampleTime) override;
 
+  std::string object_prim_path() const override;
+
   void read_matrix(float r_mat[4][4], float time, float scale, bool *r_is_constant) const;
 
   bool use_parent_xform() const
@@ -68,6 +70,9 @@ class USDXformReader : public USDPrimReader {
    *           is constant over time.
    */
   virtual std::optional<XformResult> get_local_usd_xform(float time) const;
+
+ private:
+  pxr::UsdGeomXformable get_xformable() const;
 };
 
 }  // namespace blender::io::usd

tests/python/bl_usd_import_test.py

@@ -18,12 +18,15 @@ class AbstractUSDTest(unittest.TestCase):
     @classmethod
     def setUpClass(cls):
         cls.testdir = args.testdir
-        cls._tempdir = tempfile.TemporaryDirectory()
-        cls.tempdir = pathlib.Path(cls._tempdir.name)
 
     def setUp(self):
+        self._tempdir = tempfile.TemporaryDirectory()
+        self.tempdir = pathlib.Path(self._tempdir.name)
+
         self.assertTrue(self.testdir.exists(),
                         'Test dir {0} should exist'.format(self.testdir))
+        self.assertTrue(self.tempdir.exists(),
+                        'Temp dir {0} should exist'.format(self.tempdir))
 
         # Make sure we always start with a known-empty file.
         bpy.ops.wm.open_mainfile(filepath=str(self.testdir / "empty.blend"))
@@ -1550,6 +1553,53 @@ class USDImportTest(AbstractUSDTest):
         check_image("color_121212.hdr", 1, 4, False)
         check_materials()
 
+    def test_get_prim_map_parent_xform_not_merged(self):
+        bpy.utils.register_class(GetPrimMapUsdImportHook)
+        bpy.ops.wm.usd_import(filepath=str(self.testdir / "usd_name_property_template.usda"), merge_parent_xform=False)
+        prim_map = GetPrimMapUsdImportHook.prim_map
+        bpy.utils.unregister_class(GetPrimMapUsdImportHook)
+
+        expected_prim_map = {
+            "/Cube": [bpy.data.objects["Cube.002"], bpy.data.meshes["Cube.002"]],
+            "/XformThenCube": [bpy.data.objects["XformThenCube"]],
+            "/XformThenCube/Cube": [bpy.data.objects["Cube"], bpy.data.meshes["Cube"]],
+            "/XformThenXformCube": [bpy.data.objects["XformThenXformCube"]],
+            "/XformThenXformCube/XformIntermediate": [bpy.data.objects["XformIntermediate"]],
+            "/XformThenXformCube/XformIntermediate/Cube": [bpy.data.objects["Cube.001"], bpy.data.meshes["Cube.001"]],
+            "/Material": [bpy.data.materials["Material"]],
+        }
+
+        self.assertDictEqual(prim_map, expected_prim_map)
+
+    def test_get_prim_map_parent_xform_merged(self):
+        bpy.utils.register_class(GetPrimMapUsdImportHook)
+        bpy.ops.wm.usd_import(filepath=str(self.testdir / "usd_name_property_template.usda"), merge_parent_xform=True)
+        prim_map = GetPrimMapUsdImportHook.prim_map
+        bpy.utils.unregister_class(GetPrimMapUsdImportHook)
+
+        expected_prim_map = {
+            "/Cube": [bpy.data.objects["Cube.002"], bpy.data.meshes["Cube.002"]],
+            "/XformThenCube": [bpy.data.objects["Cube"]],
+            "/XformThenCube/Cube": [bpy.data.meshes["Cube"]],
+            "/XformThenXformCube": [bpy.data.objects["XformThenXformCube"]],
+            "/XformThenXformCube/XformIntermediate": [bpy.data.objects["Cube.001"]],
+            "/XformThenXformCube/XformIntermediate/Cube": [bpy.data.meshes["Cube.001"]],
+            "/Material": [bpy.data.materials["Material"]],
+        }
+
+        self.assertDictEqual(prim_map, expected_prim_map)
+
+
+class GetPrimMapUsdImportHook(bpy.types.USDHook):
+    bl_idname = "get_prim_map_usd_import_hook"
+    bl_label = "Get Prim Map Usd Import Hook"
+
+    prim_map = None
+
+    @staticmethod
+    def on_import(context):
+        GetPrimMapUsdImportHook.prim_map = context.get_prim_map()
+
 
 def main():
     global args