From 74512cc5cb93e63d5caaf400d0c5d2ccd1753e63 Mon Sep 17 00:00:00 2001 From: Michael Kowalski Date: Fri, 13 Dec 2024 16:36:22 +0100 Subject: [PATCH] USD: on_material_import() and texture IO hooks Supporting a new on_material_import() USDHook callback. Also added support for import_texture() and export_texture() utility functions which can be called from on_material_import() and on_material_export() Python implementations, respectively. Pull Request: https://projects.blender.org/blender/blender/pulls/131559 --- .../blender/io/usd/intern/usd_capi_import.cc | 8 +- source/blender/io/usd/intern/usd_hook.cc | 246 +++++++++++++++++- source/blender/io/usd/intern/usd_hook.hh | 19 +- .../io/usd/intern/usd_reader_material.cc | 35 ++- .../io/usd/intern/usd_reader_material.hh | 7 +- .../blender/io/usd/intern/usd_reader_mesh.cc | 34 ++- .../blender/io/usd/intern/usd_reader_prim.hh | 23 +- .../blender/io/usd/intern/usd_reader_stage.cc | 61 ++++- .../blender/io/usd/intern/usd_reader_stage.hh | 13 + .../io/usd/intern/usd_writer_material.cc | 11 +- .../io/usd/intern/usd_writer_material.hh | 5 + tests/python/bl_usd_export_test.py | 110 ++++++++ tests/python/bl_usd_import_test.py | 162 +++++++++++- 13 files changed, 674 insertions(+), 60 deletions(-) diff --git a/source/blender/io/usd/intern/usd_capi_import.cc b/source/blender/io/usd/intern/usd_capi_import.cc index f26c5cb06b8..e1c97021de3 100644 --- a/source/blender/io/usd/intern/usd_capi_import.cc +++ b/source/blender/io/usd/intern/usd_capi_import.cc @@ -301,6 +301,11 @@ static void import_startjob(void *customdata, wmJobWorkerStatus *worker_status) USDStageReader *archive = new USDStageReader(stage, data->params, data->settings); + /* Ensure Python types for invoking hooks are registered. */ + register_hook_converters(); + + archive->find_material_import_hook_sources(); + data->archive = archive; archive->collect_readers(); @@ -476,8 +481,7 @@ static void import_endjob(void *customdata) data->archive->fake_users_for_unused_materials(); } - /* Ensure Python types for invoking hooks are registered. */ - register_hook_converters(); + data->archive->call_material_import_hooks(data->bmain); call_import_hooks(data->archive->stage(), data->prim_map, data->params.worker_status->reports); diff --git a/source/blender/io/usd/intern/usd_hook.cc b/source/blender/io/usd/intern/usd_hook.cc index 0fd4cb513d2..9607541eae2 100644 --- a/source/blender/io/usd/intern/usd_hook.cc +++ b/source/blender/io/usd/intern/usd_hook.cc @@ -3,7 +3,10 @@ * SPDX-License-Identifier: GPL-2.0-or-later */ #include "usd_hook.hh" + #include "usd.hh" +#include "usd_asset_utils.hh" +#include "usd_writer_material.hh" #include "BLI_utildefines.h" @@ -29,6 +32,7 @@ # include # include # include +# include # define PYTHON_NS pxr::pxr_boost::python # define REF pxr::pxr_boost::python::ref @@ -39,6 +43,7 @@ using namespace pxr::pxr_boost; # include # include # include +# include # define PYTHON_NS boost::python # define REF boost::ref @@ -171,16 +176,110 @@ struct USDSceneImportContext { /* Encapsulate arguments for material export. */ struct USDMaterialExportContext { - USDMaterialExportContext() = default; + USDMaterialExportContext() : reports(nullptr) {} - USDMaterialExportContext(pxr::UsdStageRefPtr in_stage) : stage(in_stage) {} + USDMaterialExportContext(pxr::UsdStageRefPtr in_stage, + const USDExportParams &in_params, + ReportList *in_reports) + : stage(in_stage), params(in_params), reports(in_reports) + { + } pxr::UsdStageRefPtr get_stage() const { return stage; } + /** + * Returns the USD asset export path for the given texture image. The image will be copied + * to the export directory if exporting textures is enabled in the export options. The + * function may return an empty string in case of an error. + */ + std::string export_texture(PYTHON_NS::object obj) + { + ID *id; + if (!pyrna_id_FromPyObject(obj.ptr(), &id)) { + return ""; + } + + if (!id) { + return ""; + } + + if (GS(id->name) != ID_IM) { + return ""; + } + + Image *ima = reinterpret_cast(id); + + std::string asset_path = get_tex_image_asset_filepath(ima, stage, params); + + if (params.export_textures) { + blender::io::usd::export_texture(ima, stage, params.overwrite_textures, reports); + } + + return asset_path; + } + pxr::UsdStageRefPtr stage; + USDExportParams params; + ReportList *reports; +}; + +/* Encapsulate arguments for material import. */ +struct USDMaterialImportContext { + USDMaterialImportContext() : reports(nullptr) {} + + USDMaterialImportContext(pxr::UsdStageRefPtr in_stage, + const USDImportParams &in_params, + ReportList *in_reports) + : stage(in_stage), params(in_params), reports(in_reports) + { + } + + pxr::UsdStageRefPtr get_stage() const + { + return stage; + } + + /** + * If the given texture asset path is a URI or is relative to a USDZ arhive, attempt to copy the + * texture to the local file system and returns a tuple[str, bool], containing the asset's local + * path and a boolean indicating whether the path references a temporary file (in the case where + * imported textures should be packed). The original asset path will be returned unchanged if + * it's alreay a local file or if it could not be copied to a local destination. + */ + PYTHON_NS::tuple import_texture(const std::string &asset_path) const + { + if (!should_import_asset(asset_path)) { + /* This path does not need to be imported, so return it unchanged. */ + return PYTHON_NS::make_tuple(asset_path, false); + } + + const char *textures_dir = params.import_textures_mode == USD_TEX_IMPORT_PACK ? + temp_textures_dir() : + params.import_textures_dir; + + const eUSDTexNameCollisionMode name_collision_mode = params.import_textures_mode == + USD_TEX_IMPORT_PACK ? + USD_TEX_NAME_COLLISION_OVERWRITE : + params.tex_name_collision_mode; + + std::string import_path = import_asset( + asset_path.c_str(), textures_dir, name_collision_mode, reports); + + if (import_path == asset_path) { + /* Path is unchanged. */ + return PYTHON_NS::make_tuple(asset_path, false); + } + + const bool is_temporary = params.import_textures_mode == USD_TEX_IMPORT_PACK; + return PYTHON_NS::make_tuple(import_path, is_temporary); + } + + pxr::UsdStageRefPtr stage; + USDImportParams params; + ReportList *reports; }; void register_hook_converters() @@ -215,12 +314,17 @@ void register_hook_converters() python::return_value_policy()); python::class_("USDMaterialExportContext") - .def("get_stage", &USDMaterialExportContext::get_stage); + .def("get_stage", &USDMaterialExportContext::get_stage) + .def("export_texture", &USDMaterialExportContext::export_texture); python::class_("USDSceneImportContext") .def("get_stage", &USDSceneImportContext::get_stage) .def("get_prim_map", &USDSceneImportContext::get_prim_map); + python::class_("USDMaterialImportContext") + .def("get_stage", &USDMaterialImportContext::get_stage) + .def("import_texture", &USDMaterialImportContext::import_texture); + PyGILState_Release(gilstate); } @@ -246,6 +350,8 @@ static void handle_python_error(USDHook *hook, ReportList *reports) */ class USDHookInvoker { public: + USDHookInvoker(ReportList *reports) : reports_(reports) {} + /* Attempt to call the function, if defined by the registered hooks. */ void call() { @@ -301,7 +407,7 @@ class USDHookInvoker { * required arguments, e.g., * * python::call_method(hook_obj, function_name(), arg1, arg2); */ - virtual void call_hook(PyObject *hook_obj) const = 0; + virtual void call_hook(PyObject *hook_obj) = 0; virtual void init_in_gil(){}; virtual void release_in_gil(){}; @@ -316,9 +422,8 @@ class OnExportInvoker : public USDHookInvoker { public: OnExportInvoker(pxr::UsdStageRefPtr stage, Depsgraph *depsgraph, ReportList *reports) - : hook_context_(stage, depsgraph) + : USDHookInvoker(reports), hook_context_(stage, depsgraph) { - reports_ = reports; } protected: @@ -327,7 +432,7 @@ class OnExportInvoker : public USDHookInvoker { return "on_export"; } - void call_hook(PyObject *hook_obj) const override + void call_hook(PyObject *hook_obj) override { python::call_method(hook_obj, function_name(), REF(hook_context_)); } @@ -343,11 +448,13 @@ class OnMaterialExportInvoker : public USDHookInvoker { OnMaterialExportInvoker(pxr::UsdStageRefPtr stage, Material *material, const pxr::UsdShadeMaterial &usd_material, + const USDExportParams &export_params, ReportList *reports) - : hook_context_(stage), usd_material_(usd_material) + : USDHookInvoker(reports), + hook_context_(stage, export_params, reports), + usd_material_(usd_material) { material_ptr_ = RNA_pointer_create(nullptr, &RNA_Material, material); - reports_ = reports; } protected: @@ -356,7 +463,7 @@ class OnMaterialExportInvoker : public USDHookInvoker { return "on_material_export"; } - void call_hook(PyObject *hook_obj) const override + void call_hook(PyObject *hook_obj) override { python::call_method( hook_obj, function_name(), REF(hook_context_), material_ptr_, usd_material_); @@ -369,9 +476,8 @@ class OnImportInvoker : public USDHookInvoker { public: OnImportInvoker(pxr::UsdStageRefPtr stage, const ImportedPrimMap &prim_map, ReportList *reports) - : hook_context_(stage, prim_map) + : USDHookInvoker(reports), hook_context_(stage, prim_map) { - reports_ = reports; } protected: @@ -380,7 +486,7 @@ class OnImportInvoker : public USDHookInvoker { return "on_import"; } - void call_hook(PyObject *hook_obj) const override + void call_hook(PyObject *hook_obj) override { python::call_method(hook_obj, function_name(), REF(hook_context_)); } @@ -391,6 +497,85 @@ class OnImportInvoker : public USDHookInvoker { } }; +class MaterialImportPollInvoker : public USDHookInvoker { + private: + USDMaterialImportContext hook_context_; + pxr::UsdShadeMaterial usd_material_; + bool result_; + + public: + MaterialImportPollInvoker(pxr::UsdStageRefPtr stage, + const pxr::UsdShadeMaterial &usd_material, + const USDImportParams &import_params, + ReportList *reports) + : USDHookInvoker(reports), + hook_context_(stage, import_params, reports), + usd_material_(usd_material), + result_(false) + { + } + + bool result() const + { + return result_; + } + + protected: + const char *function_name() const override + { + return "material_import_poll"; + } + + void call_hook(PyObject *hook_obj) override + { + // If we already know that one of the registered hook classes can import the material + // because it returned true in a previous invocation of the callback, we skip the call. + if (!result_) { + result_ = python::call_method( + hook_obj, function_name(), REF(hook_context_), usd_material_); + } + } +}; + +class OnMaterialImportInvoker : public USDHookInvoker { + private: + USDMaterialImportContext hook_context_; + pxr::UsdShadeMaterial usd_material_; + PointerRNA material_ptr_; + bool result_; + + public: + OnMaterialImportInvoker(pxr::UsdStageRefPtr stage, + Material *material, + const pxr::UsdShadeMaterial &usd_material, + const USDImportParams &import_params, + ReportList *reports) + : USDHookInvoker(reports), + hook_context_(stage, import_params, reports), + usd_material_(usd_material), + result_(false) + { + material_ptr_ = RNA_pointer_create(nullptr, &RNA_Material, material); + } + + bool result() const + { + return result_; + } + + protected: + const char *function_name() const override + { + return "on_material_import"; + } + + void call_hook(PyObject *hook_obj) override + { + result_ |= python::call_method( + hook_obj, function_name(), REF(hook_context_), material_ptr_, usd_material_); + } +}; + void call_export_hooks(pxr::UsdStageRefPtr stage, Depsgraph *depsgraph, ReportList *reports) { if (hook_list().empty()) { @@ -404,13 +589,15 @@ void call_export_hooks(pxr::UsdStageRefPtr stage, Depsgraph *depsgraph, ReportLi void call_material_export_hooks(pxr::UsdStageRefPtr stage, Material *material, const pxr::UsdShadeMaterial &usd_material, + const USDExportParams &export_params, ReportList *reports) { if (hook_list().empty()) { return; } - OnMaterialExportInvoker on_material_export(stage, material, usd_material, reports); + OnMaterialExportInvoker on_material_export( + stage, material, usd_material, export_params, reports); on_material_export.call(); } @@ -426,6 +613,37 @@ void call_import_hooks(pxr::UsdStageRefPtr stage, on_import.call(); } +bool have_material_import_hook(pxr::UsdStageRefPtr stage, + const pxr::UsdShadeMaterial &usd_material, + const USDImportParams &import_params, + ReportList *reports) +{ + if (hook_list().empty()) { + return false; + } + + MaterialImportPollInvoker poll(stage, usd_material, import_params, reports); + poll.call(); + + return poll.result(); +} + +bool call_material_import_hooks(pxr::UsdStageRefPtr stage, + Material *material, + const pxr::UsdShadeMaterial &usd_material, + const USDImportParams &import_params, + ReportList *reports) +{ + if (hook_list().empty()) { + return false; + } + + OnMaterialImportInvoker on_material_import( + stage, material, usd_material, import_params, reports); + on_material_import.call(); + return on_material_import.result(); +} + } // namespace blender::io::usd #undef REF diff --git a/source/blender/io/usd/intern/usd_hook.hh b/source/blender/io/usd/intern/usd_hook.hh index 9b2734c3194..b3b85fce81f 100644 --- a/source/blender/io/usd/intern/usd_hook.hh +++ b/source/blender/io/usd/intern/usd_hook.hh @@ -17,19 +17,22 @@ struct ReportList; namespace blender::io::usd { +struct USDExportParams; +struct USDImportParams; using ImportedPrimMap = Map>; /** Ensure classes and type converters necessary for invoking import and export hooks * are registered. */ void register_hook_converters(); -/** Call the 'on_export' chaser function defined in the registered USDHook classes. */ +/** Call the 'on_export' chaser function defined in the registered #USDHook classes. */ void call_export_hooks(pxr::UsdStageRefPtr stage, Depsgraph *depsgraph, ReportList *reports); /** Call the 'on_material_export' hook functions defined in the registered #USDHook classes. */ void call_material_export_hooks(pxr::UsdStageRefPtr stage, Material *material, const pxr::UsdShadeMaterial &usd_material, + const USDExportParams &export_params, ReportList *reports); /** Call the 'on_import' chaser function defined in the registered USDHook classes. */ @@ -37,4 +40,18 @@ void call_import_hooks(pxr::UsdStageRefPtr stage, const ImportedPrimMap &imported_id_links, ReportList *reports); +/** Returns true if there is a registered #USDHook class that can convert the given material. */ +bool have_material_import_hook(pxr::UsdStageRefPtr stage, + const pxr::UsdShadeMaterial &usd_material, + const USDImportParams &import_params, + ReportList *reports); + +/** Call the 'on_material_import' hook functions defined in the registered #USDHook classes. + * Returns true if any of the hooks were successful, false otherwise. */ +bool call_material_import_hooks(pxr::UsdStageRefPtr stage, + Material *material, + const pxr::UsdShadeMaterial &usd_material, + const USDImportParams &import_params, + ReportList *reports); + } // namespace blender::io::usd diff --git a/source/blender/io/usd/intern/usd_reader_material.cc b/source/blender/io/usd/intern/usd_reader_material.cc index b988b1958fd..3bfa4fa4305 100644 --- a/source/blender/io/usd/intern/usd_reader_material.cc +++ b/source/blender/io/usd/intern/usd_reader_material.cc @@ -455,7 +455,8 @@ USDMaterialReader::USDMaterialReader(const USDImportParams ¶ms, Main *bmain) { } -Material *USDMaterialReader::add_material(const pxr::UsdShadeMaterial &usd_material) const +Material *USDMaterialReader::add_material(const pxr::UsdShadeMaterial &usd_material, + const bool read_usd_preview) const { if (!(bmain_ && usd_material)) { return nullptr; @@ -467,17 +468,8 @@ Material *USDMaterialReader::add_material(const pxr::UsdShadeMaterial &usd_mater Material *mtl = BKE_material_add(bmain_, mtl_name.c_str()); id_us_min(&mtl->id); - /* Get the UsdPreviewSurface shader source for the material, - * if there is one. */ - pxr::UsdShadeShader usd_preview; - if (get_usd_preview_surface(usd_material, usd_preview)) { - - set_viewport_material_props(mtl, usd_preview); - - /* Optionally, create shader nodes to represent a UsdPreviewSurface. */ - if (params_.import_usd_preview) { - import_usd_preview(mtl, usd_preview); - } + if (read_usd_preview) { + import_usd_preview(mtl, usd_material); } /* Load custom properties directly from the Material's prim. */ @@ -487,7 +479,24 @@ Material *USDMaterialReader::add_material(const pxr::UsdShadeMaterial &usd_mater } void USDMaterialReader::import_usd_preview(Material *mtl, - const pxr::UsdShadeShader &usd_shader) const + const pxr::UsdShadeMaterial &usd_material) const +{ + /* Get the UsdPreviewSurface shader source for the material, + * if there is one. */ + pxr::UsdShadeShader usd_preview; + if (get_usd_preview_surface(usd_material, usd_preview)) { + + set_viewport_material_props(mtl, usd_preview); + + /* Optionally, create shader nodes to represent a UsdPreviewSurface. */ + if (params_.import_usd_preview) { + import_usd_preview_nodes(mtl, usd_preview); + } + } +} + +void USDMaterialReader::import_usd_preview_nodes(Material *mtl, + const pxr::UsdShadeShader &usd_shader) const { if (!(bmain_ && mtl && usd_shader)) { return; diff --git a/source/blender/io/usd/intern/usd_reader_material.hh b/source/blender/io/usd/intern/usd_reader_material.hh index 7133cc5350d..08f5ddefa9d 100644 --- a/source/blender/io/usd/intern/usd_reader_material.hh +++ b/source/blender/io/usd/intern/usd_reader_material.hh @@ -96,7 +96,10 @@ class USDMaterialReader { public: USDMaterialReader(const USDImportParams ¶ms, Main *bmain); - Material *add_material(const pxr::UsdShadeMaterial &usd_material) const; + Material *add_material(const pxr::UsdShadeMaterial &usd_material, + bool read_usd_preview = true) const; + + void import_usd_preview(Material *mtl, const pxr::UsdShadeMaterial &usd_material) const; /** Get the wmJobWorkerStatus-provided `reports` list pointer, to use with the BKE_report API. */ ReportList *reports() const @@ -106,7 +109,7 @@ class USDMaterialReader { protected: /** Create the Principled BSDF shader node network. */ - void import_usd_preview(Material *mtl, const pxr::UsdShadeShader &usd_shader) const; + void import_usd_preview_nodes(Material *mtl, const pxr::UsdShadeShader &usd_shader) const; void set_principled_node_inputs(bNode *principled_node, bNodeTree *ntree, diff --git a/source/blender/io/usd/intern/usd_reader_mesh.cc b/source/blender/io/usd/intern/usd_reader_mesh.cc index cdd1a71fa5a..99b153857e0 100644 --- a/source/blender/io/usd/intern/usd_reader_mesh.cc +++ b/source/blender/io/usd/intern/usd_reader_mesh.cc @@ -9,6 +9,7 @@ #include "usd.hh" #include "usd_attribute_utils.hh" #include "usd_hash_types.hh" +#include "usd_hook.hh" #include "usd_mesh_utils.hh" #include "usd_reader_material.hh" #include "usd_skel_convert.hh" @@ -92,8 +93,7 @@ static void assign_materials(Main *bmain, const blender::Map &mat_index_map, const blender::io::usd::USDImportParams ¶ms, pxr::UsdStageRefPtr stage, - blender::Map &mat_name_to_mat, - blender::Map &usd_path_to_mat_name) + const blender::io::usd::ImportSettings &settings) { using namespace blender::io::usd; if (!(stage && bmain && ob)) { @@ -108,7 +108,7 @@ static void assign_materials(Main *bmain, for (const auto item : mat_index_map.items()) { Material *assigned_mat = blender::io::usd::find_existing_material( - item.key, params, mat_name_to_mat, usd_path_to_mat_name); + item.key, params, settings.mat_name_to_mat, settings.usd_path_to_mat_name); if (!assigned_mat) { /* Blender material doesn't exist, so create it now. */ @@ -122,8 +122,12 @@ static void assign_materials(Main *bmain, continue; } - /* Add the Blender material. */ - assigned_mat = mat_reader.add_material(usd_mat); + const bool have_import_hook = settings.mat_import_hook_sources.contains( + item.key.GetAsString()); + + /* Add the Blender material. If we have an import hook which can handle this material + * we don't import USD Preview Surface shaders. */ + assigned_mat = mat_reader.add_material(usd_mat, !have_import_hook); if (!assigned_mat) { CLOG_WARN(&LOG, @@ -133,12 +137,19 @@ static void assign_materials(Main *bmain, } const std::string mat_name = make_safe_name(assigned_mat->id.name + 2, true); - mat_name_to_mat.lookup_or_add_default(mat_name) = assigned_mat; + settings.mat_name_to_mat.lookup_or_add_default(mat_name) = assigned_mat; if (params.mtl_name_collision_mode == USD_MTL_NAME_COLLISION_MAKE_UNIQUE) { /* Record the name of the Blender material we created for the USD material * with the given path. */ - usd_path_to_mat_name.lookup_or_add_default(item.key.GetAsString()) = mat_name; + settings.usd_path_to_mat_name.lookup_or_add_default(item.key.GetAsString()) = mat_name; + } + + if (have_import_hook) { + /* Defer invoking the hook to convert the material till we can do so from + * the main thread. */ + settings.usd_path_to_mat_for_hook.lookup_or_add_default( + item.key.GetAsString()) = assigned_mat; } } @@ -736,13 +747,8 @@ void USDMeshReader::readFaceSetsSample(Main *bmain, Mesh *mesh, const double mot if (this->settings_->mat_name_to_mat.is_empty()) { build_material_map(bmain, &this->settings_->mat_name_to_mat); } - utils::assign_materials(bmain, - object_, - mat_map, - this->import_params_, - this->prim_.GetStage(), - this->settings_->mat_name_to_mat, - this->settings_->usd_path_to_mat_name); + utils::assign_materials( + bmain, object_, mat_map, this->import_params_, this->prim_.GetStage(), *this->settings_); } Mesh *USDMeshReader::read_mesh(Mesh *existing_mesh, diff --git a/source/blender/io/usd/intern/usd_reader_prim.hh b/source/blender/io/usd/intern/usd_reader_prim.hh index 1f3a3fb9804..c3a3ab71774 100644 --- a/source/blender/io/usd/intern/usd_reader_prim.hh +++ b/source/blender/io/usd/intern/usd_reader_prim.hh @@ -10,6 +10,7 @@ #include "usd.hh" #include "BLI_map.hh" +#include "BLI_set.hh" #include "WM_types.hh" @@ -35,16 +36,26 @@ struct ImportSettings { std::function get_cache_file{}; + /* + * The fields below are mutable because they are used to keep track + * of what the importer is doing. This is necessary even when all + * the other import settings are to remain const. + */ + /* Map a USD material prim path to a Blender material name. - * This map is updated by readers during stage traversal. - * This field is mutable because it is used to keep track - * of what the importer is doing. This is necessary even - * when all the other import settings are to remain const. */ + * This map is updated by readers during stage traversal. */ mutable blender::Map usd_path_to_mat_name{}; /* Map a material name to Blender material. - * This map is updated by readers during stage traversal, - * and is mutable similar to the map above. */ + * This map is updated by readers during stage traversal. */ mutable blender::Map mat_name_to_mat{}; + /* Map a USD material prim path to a Blender material to be + * converted by invoking the 'on_material_import' USD hook. + * This map is updated by readers during stage traversal. */ + mutable blender::Map usd_path_to_mat_for_hook{}; + /* Set of paths to USD material prims that can be converted by the + * 'on_material_import' USD hook. For efficiency this set should + * be populated prior to stage traversal. */ + mutable blender::Set mat_import_hook_sources{}; /* We use the stage metersPerUnit to convert camera properties from USD scene units to the * correct millimeter scale that Blender uses for camera parameters. */ diff --git a/source/blender/io/usd/intern/usd_reader_stage.cc b/source/blender/io/usd/intern/usd_reader_stage.cc index a9acaf5eb33..3d09b315834 100644 --- a/source/blender/io/usd/intern/usd_reader_stage.cc +++ b/source/blender/io/usd/intern/usd_reader_stage.cc @@ -3,6 +3,8 @@ * SPDX-License-Identifier: GPL-2.0-or-later */ #include "usd_reader_stage.hh" + +#include "usd_hook.hh" #include "usd_reader_camera.hh" #include "usd_reader_curve.hh" #include "usd_reader_instance.hh" @@ -20,6 +22,7 @@ #include "usd_utils.hh" #include +#include #include #include #include @@ -535,8 +538,12 @@ void USDStageReader::import_all_materials(Main *bmain) continue; } - /* Add the material now. */ - Material *new_mtl = mtl_reader.add_material(usd_mtl); + /* Can the material be handled by an iport hook? */ + const bool have_import_hook = settings_.mat_import_hook_sources.contains(mtl_path); + + /* Add the Blender material. If we have an import hook which can handle this material + * we don't import USD Preview Surface shaders. */ + Material *new_mtl = mtl_reader.add_material(usd_mtl, !have_import_hook); BLI_assert_msg(new_mtl, "Failed to create material"); const std::string mtl_name = make_safe_name(new_mtl->id.name + 2, true); @@ -549,6 +556,12 @@ void USDStageReader::import_all_materials(Main *bmain) settings_.usd_path_to_mat_name.lookup_or_add_default( prim.GetPath().GetAsString()) = mtl_name; } + + if (have_import_hook) { + /* Defer invoking the hook to convert the material till we can do so from + * the main thread. */ + settings_.usd_path_to_mat_for_hook.lookup_or_add_default(mtl_path) = new_mtl; + } } } @@ -568,6 +581,50 @@ void USDStageReader::fake_users_for_unused_materials() } } +void USDStageReader::find_material_import_hook_sources() +{ + pxr::UsdPrimRange range = stage_->Traverse(); + for (pxr::UsdPrim prim : range) { + if (prim.IsA()) { + pxr::UsdShadeMaterial usd_mat(prim); + if (have_material_import_hook(stage_, usd_mat, params_, reports())) { + settings_.mat_import_hook_sources.add(prim.GetPath().GetAsString()); + } + } + } +} + +void USDStageReader::call_material_import_hooks(Main *bmain) const +{ + if (settings_.usd_path_to_mat_for_hook.is_empty()) { + /* No materials can be converted by a hook. */ + return; + } + + for (const auto item : settings_.usd_path_to_mat_for_hook.items()) { + pxr::UsdPrim prim = stage_->GetPrimAtPath(pxr::SdfPath(item.key)); + + pxr::UsdShadeMaterial usd_mtl(prim); + if (!usd_mtl) { + continue; + } + + bool success = blender::io::usd::call_material_import_hooks( + stage_, item.value, usd_mtl, params_, reports()); + + if (!success) { + /* None of the hooks succeeded, so fall back on importing USD Preview Surface if possible. */ + CLOG_WARN(&LOG, + "USD hook 'on_material_import' for material %s failed, attempting to convert USD " + "Preview Surface material", + usd_mtl.GetPath().GetAsString().c_str()); + + USDMaterialReader mat_reader(this->params_, bmain); + mat_reader.import_usd_preview(item.value, usd_mtl); + } + } +} + void USDStageReader::clear_readers() { for (USDPrimReader *reader : readers_) { diff --git a/source/blender/io/usd/intern/usd_reader_stage.hh b/source/blender/io/usd/intern/usd_reader_stage.hh index 4170e05b859..29e4e188947 100644 --- a/source/blender/io/usd/intern/usd_reader_stage.hh +++ b/source/blender/io/usd/intern/usd_reader_stage.hh @@ -87,6 +87,19 @@ class USDStageReader { * materials. */ void fake_users_for_unused_materials(); + /** + * Discover the USD materials that can be converted + * by material import hook add-ons. + */ + void find_material_import_hook_sources(); + + /** + * Invoke USD hook add-ons to convert materials. This function + * should be called from the main thread and not from a + * background job. + */ + void call_material_import_hooks(struct Main *bmain) const; + bool valid() const; pxr::UsdStageRefPtr stage() diff --git a/source/blender/io/usd/intern/usd_writer_material.cc b/source/blender/io/usd/intern/usd_writer_material.cc index 9e5d83a3940..ce0d5aa6851 100644 --- a/source/blender/io/usd/intern/usd_writer_material.cc +++ b/source/blender/io/usd/intern/usd_writer_material.cc @@ -1287,10 +1287,10 @@ static void copy_single_file(const Image *ima, } } -static void export_texture(Image *ima, - const pxr::UsdStageRefPtr stage, - const bool allow_overwrite, - ReportList *reports) +void export_texture(Image *ima, + const pxr::UsdStageRefPtr stage, + const bool allow_overwrite, + ReportList *reports) { std::string dest_dir = get_export_textures_dir(stage); if (dest_dir.empty()) { @@ -1618,7 +1618,8 @@ pxr::UsdShadeMaterial create_usd_material(const USDExporterContext &usd_export_c } #endif - call_material_export_hooks(usd_export_context.stage, material, usd_material, reports); + call_material_export_hooks( + usd_export_context.stage, material, usd_material, usd_export_context.export_params, reports); return usd_material; } diff --git a/source/blender/io/usd/intern/usd_writer_material.hh b/source/blender/io/usd/intern/usd_writer_material.hh index 3d97dc5a963..895b3b98e2e 100644 --- a/source/blender/io/usd/intern/usd_writer_material.hh +++ b/source/blender/io/usd/intern/usd_writer_material.hh @@ -40,6 +40,11 @@ void export_texture(bNode *node, const bool allow_overwrite = false, ReportList *reports = nullptr); +void export_texture(Image *ima, + const pxr::UsdStageRefPtr stage, + const bool allow_overwrite = false, + ReportList *reports = nullptr); + /** * Gets an asset path for the given texture image / node. The resulting path * may be absolute, relative to the USD file, or in a 'textures' directory diff --git a/tests/python/bl_usd_export_test.py b/tests/python/bl_usd_export_test.py index 44f8af88f38..347b2860185 100644 --- a/tests/python/bl_usd_export_test.py +++ b/tests/python/bl_usd_export_test.py @@ -1124,6 +1124,86 @@ class USDExportTest(AbstractUSDTest): self.assertTupleEqual(expected, actual) + def test_texture_export_hook(self): + """Exporting textures from on_material_export USD hook.""" + + # Clear USD hook results. + ExportTextureUSDHook.exported_textures = {} + + bpy.utils.register_class(ExportTextureUSDHook) + bpy.ops.wm.open_mainfile(filepath=str(self.testdir / "usd_materials_export.blend")) + + export_path = self.tempdir / "usd_materials_export.usda" + + self.export_and_validate( + filepath=str(export_path), + export_materials=True, + generate_preview_surface=False, + ) + + # Verify that the exported texture paths were returned as expected. + expected = {'/root/_materials/Transforms': './textures/test_grid_.png', + '/root/_materials/Clip_With_Round': './textures/test_grid_.png', + '/root/_materials/NormalMap': './textures/test_normal.exr', + '/root/_materials/Material': './textures/test_grid_.png', + '/root/_materials/Clip_With_LessThanInvert': './textures/test_grid_.png', + '/root/_materials/NormalMap_Scale_Bias': './textures/test_normal_invertY.exr'} + + self.assertDictEqual(ExportTextureUSDHook.exported_textures, + expected, + "Unexpected texture export paths") + + bpy.utils.unregister_class(ExportTextureUSDHook) + + # Verify that the texture files were copied as expected. + tex_names = ['test_grid_1001.png', 'test_grid_1002.png', + 'test_normal.exr', 'test_normal_invertY.exr'] + + for name in tex_names: + tex_path = self.tempdir / "textures" / name + self.assertTrue(tex_path.exists(), + f"Exported texture {tex_path} doesn't exist") + + def test_inmem_pack_texture_export_hook(self): + """Exporting packed and in memory textures from on_material_export USD hook.""" + + # Clear hook results. + ExportTextureUSDHook.exported_textures = {} + + bpy.utils.register_class(ExportTextureUSDHook) + bpy.ops.wm.open_mainfile(filepath=str(self.testdir / "usd_materials_inmem_pack.blend")) + + export_path = self.tempdir / "usd_materials_inmem_pack.usda" + + self.export_and_validate( + filepath=str(export_path), + export_materials=True, + generate_preview_surface=False, + ) + + # Verify that the exported texture paths were returned as expected. + expected = {'/root/_materials/MAT_pack_udim': './textures/test_grid_.png', + '/root/_materials/MAT_pack_single': './textures/test_single.png', + '/root/_materials/MAT_inmem_udim': './textures/inmem_udim..png', + '/root/_materials/MAT_inmem_single': './textures/inmem_single.png'} + + self.assertDictEqual(ExportTextureUSDHook.exported_textures, + expected, + "Unexpected texture export paths") + + bpy.utils.unregister_class(ExportTextureUSDHook) + + # Verify that the texture files were copied as expected. + tex_names = ['test_grid_1001.png', 'test_grid_1002.png', + 'test_single.png', + 'inmem_udim.1001.png', 'inmem_udim.1002.png', + 'inmem_single.png'] + + for name in tex_names: + tex_path = self.tempdir / "textures" / name + self.assertTrue(tex_path.exists(), + f"Exported texture {tex_path} doesn't exist") + class USDHookBase(): instructions = {} @@ -1208,6 +1288,36 @@ class USDHook2(USDHookBase, bpy.types.USDHook): return USDHookBase.do_on_import(USDHook2.bl_label, import_context) +class ExportTextureUSDHook(bpy.types.USDHook): + bl_idname = "export_texture_usd_hook" + bl_label = "Export Texture USD Hook" + + exported_textures = {} + + @staticmethod + def on_material_export(export_context, bl_material, usd_material): + """ + If a texture image node exists in the given material's + node tree, call exprt_texture() on the image and cache + the returned path. + """ + tex_image_node = None + if bl_material and bl_material.node_tree: + for node in bl_material.node_tree.nodes: + if node.type == 'TEX_IMAGE': + tex_image_node = node + + if tex_image_node is None: + return False + + tex_path = export_context.export_texture(tex_image_node.image) + + ExportTextureUSDHook.exported_textures[usd_material.GetPath() + .pathString] = tex_path + + return True + + def main(): global args import argparse diff --git a/tests/python/bl_usd_import_test.py b/tests/python/bl_usd_import_test.py index 132f42f1bb3..6591f55abc6 100644 --- a/tests/python/bl_usd_import_test.py +++ b/tests/python/bl_usd_import_test.py @@ -7,7 +7,7 @@ import pathlib import sys import tempfile import unittest -from pxr import Gf, Sdf, Usd, UsdGeom, UsdShade +from pxr import Ar, Gf, Sdf, Usd, UsdGeom, UsdShade import bpy @@ -1589,6 +1589,73 @@ class USDImportTest(AbstractUSDTest): self.assertDictEqual(prim_map, expected_prim_map) + def test_material_import_usd_hook(self): + """Test importing color from an mtlx shader.""" + + bpy.utils.register_class(ImportMtlxColorUSDHook) + bpy.ops.wm.usd_import(filepath=str(self.testdir / "usd_simple_mtlx.usda")) + bpy.utils.unregister_class(ImportMtlxColorUSDHook) + + # Check that the correct color was read. + imported_color = ImportMtlxColorUSDHook.imported_color + self.assertEqual(imported_color, [0, 1, 0, 1], "Wrong mtlx shader color imported") + + # Check that a Principled BSDF shader with the expected Base Color input. + # was created. + mtl = bpy.data.materials["Material"] + self.assertTrue(mtl.use_nodes) + bsdf = mtl.node_tree.nodes.get("Principled BSDF") + self.assertIsNotNone(bsdf) + base_color_input = bsdf.inputs['Base Color'] + self.assertEqual(self.round_vector(base_color_input.default_value), [0, 1, 0, 1]) + + def test_usd_hook_import_texture(self): + """ + Test importing a texture from a USDZ archive, using two different + texture import modes. + """ + + bpy.utils.register_class(ImportMtlxTextureUSDHook) + bpy.ops.wm.usd_import(filepath=str(self.testdir / "usd_simple_mtlx_texture.usdz"), + import_textures_mode='IMPORT_COPY', + import_textures_dir=str(self.tempdir / "textures")) + + # The resolved path should be a package-relative path for the USDZ, in this case + # self.testdir / "usd_simple_mtlx_texture.usdz[textures/test_single.png]" + resolved_path = ImportMtlxTextureUSDHook.resolved_path + + self.assertTrue(Ar.IsPackageRelativePath(resolved_path), + "Resolved path is not relative to the USDZ") + path_inner = Ar.SplitPackageRelativePathInner(resolved_path) + self.assertEqual(path_inner[1], "textures/test_single.png", + "Resolved path does not have the expected format") + + # Confirm that file was copied from the USDZ archive to the textures + # directory. + import_path = ImportMtlxTextureUSDHook.result[0] + self.assertTrue(pathlib.Path(import_path).exists(), + "Imported texture does not exist") + # Path should not be temporary + is_temporary = ImportMtlxTextureUSDHook.result[1] + self.assertFalse(is_temporary, + "Imported texture should not be temporary") + + # Repeat the test with texture packing enabled. + bpy.ops.wm.usd_import(filepath=str(self.testdir / "usd_simple_mtlx_texture.usdz"), + import_textures_mode='IMPORT_PACK', + import_textures_dir="") + + # Confirm that the copied file exists. + import_path = ImportMtlxTextureUSDHook.result[0] + self.assertTrue(pathlib.Path(import_path).exists(), + "Imported texture does not exist") + # Path should be temporary + is_temporary = ImportMtlxTextureUSDHook.result[1] + self.assertTrue(is_temporary, + "Imported texture should be temporary") + + bpy.utils.unregister_class(ImportMtlxTextureUSDHook) + class GetPrimMapUsdImportHook(bpy.types.USDHook): bl_idname = "get_prim_map_usd_import_hook" @@ -1601,6 +1668,99 @@ class GetPrimMapUsdImportHook(bpy.types.USDHook): GetPrimMapUsdImportHook.prim_map = context.get_prim_map() +class ImportMtlxColorUSDHook(bpy.types.USDHook): + """ + Simple test for importing an mtxl shader from the usd_simple_mtlx.usda + test file. + """ + bl_idname = "import_mtlx_color_usd_hook" + bl_label = "Import Mtlx Color USD Hook" + + imported_color = None + + @staticmethod + def material_import_poll(import_context, usd_material): + # We can import the material if it has an 'mtlx' context. + surf_output = usd_material.GetSurfaceOutput("mtlx") + return bool(surf_output) + + @staticmethod + def on_material_import(import_context, bl_material, usd_material): + + # Get base_color from connected surface shader. + surf_output = usd_material.GetSurfaceOutput("mtlx") + assert surf_output + source = surf_output.GetConnectedSource() + assert source + shader = UsdShade.Shader(source[0]) + assert shader + color_attr = shader.GetInput("base_color") + assert color_attr + # Get the authored default color. + color = color_attr.Get() + + # Add a Principled BSDF shader and set its 'Base Color' input to + # the color we read from mtlx. + bl_material.use_nodes = True + node_tree = bl_material.node_tree + nodes = node_tree.nodes + bsdf = nodes.get("Principled BSDF") + assert bsdf + color4 = [color[0], color[1], color[2], 1] + ImportMtlxColorUSDHook.imported_color = color4 + bsdf.inputs['Base Color'].default_value = color4 + + return True + + +class ImportMtlxTextureUSDHook(bpy.types.USDHook): + """ + Simple test for importing a texture file from the + usd_simple_mtlx_texture.usdz archive test file. + """ + bl_idname = "import_mtlx_texture_usd_hook" + bl_label = "Import Mtlx Texture USD Hook" + + resolved_path = None + result = None + + @staticmethod + def material_import_poll(import_context, usd_material): + # We can import the material if it has an 'mtlx' context. + surf_output = usd_material.GetSurfaceOutput("mtlx") + return bool(surf_output) + + @staticmethod + def on_material_import(import_context, bl_material, usd_material): + # For the test, we get the texture file input of a shader known + # to exist in the usd_simple_mtlx_texture.usdz archive. + surf_output = usd_material.GetSurfaceOutput("mtlx") + assert surf_output + stage = import_context.get_stage() + assert stage + prim = stage.GetPrimAtPath("/root/_materials/Material/NodeGraphs/Image_Texture_Color") + assert prim + tex_node = UsdShade.Shader(prim) + assert tex_node + file_attr = tex_node.GetInput("file") + assert file_attr + file = file_attr.Get() + + # Record the file's resolved path, which should be a package-relative + # path, e.g., + # c:/foo/bar/usd_simple_mtlx_texture.usdz[textures/test_single.png + resolved_path = file.resolvedPath + ImportMtlxTextureUSDHook.resolved_path = resolved_path + + result = import_context.import_texture(resolved_path) + # Record the returned result tuple. The first element of the tuple + # is the texture path and the second is a flag indicating whether the + # returned path references a temporary file. + ImportMtlxTextureUSDHook.result = result + + return True + + def main(): global args import argparse