diff --git a/doc/python_api/examples/bpy.types.USDHook.py b/doc/python_api/examples/bpy.types.USDHook.py index 7f95d8000b8..a211bd0b7bf 100644 --- a/doc/python_api/examples/bpy.types.USDHook.py +++ b/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 diff --git a/source/blender/io/usd/intern/usd_capi_import.cc b/source/blender/io/usd/intern/usd_capi_import.cc index 7d7b22bc36e..3fae95d7b0a 100644 --- a/source/blender/io/usd/intern/usd_capi_import.cc +++ b/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(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 diff --git a/source/blender/io/usd/intern/usd_hook.cc b/source/blender/io/usd/intern/usd_hook.cc index ddfc44cf4f5..0ed541e78f1 100644 --- a/source/blender/io/usd/intern/usd_hook.cc +++ b/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( + (*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") - .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(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(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(); } diff --git a/source/blender/io/usd/intern/usd_hook.hh b/source/blender/io/usd/intern/usd_hook.hh index 1d412cd4981..8404257d886 100644 --- a/source/blender/io/usd/intern/usd_hook.hh +++ b/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>; + /** 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 diff --git a/source/blender/io/usd/intern/usd_reader_prim.hh b/source/blender/io/usd/intern/usd_reader_prim.hh index fa746ba50d7..1f3a3fb9804 100644 --- a/source/blender/io/usd/intern/usd_reader_prim.hh +++ b/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; diff --git a/source/blender/io/usd/intern/usd_reader_stage.hh b/source/blender/io/usd/intern/usd_reader_stage.hh index 81953c2ad54..4170e05b859 100644 --- a/source/blender/io/usd/intern/usd_reader_stage.hh +++ b/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 readers_; diff --git a/source/blender/io/usd/intern/usd_reader_xform.cc b/source/blender/io/usd/intern/usd_reader_xform.cc index ca9211d8e11..2ba2296aee5 100644 --- a/source/blender/io/usd/intern/usd_reader_xform.cc +++ b/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 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 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 diff --git a/source/blender/io/usd/intern/usd_reader_xform.hh b/source/blender/io/usd/intern/usd_reader_xform.hh index 82f8741a151..a1c000389a5 100644 --- a/source/blender/io/usd/intern/usd_reader_xform.hh +++ b/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 get_local_usd_xform(float time) const; + + private: + pxr::UsdGeomXformable get_xformable() const; }; } // namespace blender::io::usd diff --git a/tests/python/bl_usd_import_test.py b/tests/python/bl_usd_import_test.py index 3a7f30e04d5..132f42f1bb3 100644 --- a/tests/python/bl_usd_import_test.py +++ b/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