Merge pull request #1584 from zenustech/abc_save_faceset

Abc save faceset
zenustech · Dec 7, 2023 · 6c4fa7e · 6c4fa7e
2 parents 37377b5 + 98a7bc5
commit 6c4fa7e
Show file tree

Hide file tree

Showing 2 changed files with 64 additions and 1 deletion.
diff --git a/projects/Alembic/ReadAlembic.cpp b/projects/Alembic/ReadAlembic.cpp
@@ -822,5 +822,41 @@ ZENDEFNODE(AlembicSplitByName, {
     {"alembic"},
 });
 
+struct CopyPosAndNrmByIndex: INode {
+    void apply() override {
+        auto prim = get_input<PrimitiveObject>("prim");
+        auto prims = get_input<ListObject>("list")->get<PrimitiveObject>();
+        for (auto p: prims) {
+            size_t size = p->size();
+            auto index = p->attr<int>("index");
+            for (auto i = 0; i < size; i++) {
+                prim->verts[index[i]] = p->verts[i];
+            }
+            if (prim->verts.attr_is<vec3f>("nrm")) {
+                auto &nrm = prim->verts.attr<vec3f>("nrm");
+                auto &nrm_sub = p->verts.attr<vec3f>("nrm");
+                for (auto i = 0; i < size; i++) {
+                    nrm[index[i]] = nrm_sub[i];
+                }
+            }
+        }
+
+        set_output("out", prim);
+    }
+};
+
+ZENDEFNODE(CopyPosAndNrmByIndex, {
+    {
+        {"prim"},
+        {"list", "list"},
+    },
+    {
+        {"out"},
+    },
+    {},
+    {"alembic"},
+});
+
+
 
 } // namespace zeno
diff --git a/projects/Alembic/WriteAlembic.cpp b/projects/Alembic/WriteAlembic.cpp
@@ -10,6 +10,8 @@
 #include <Alembic/AbcCoreOgawa/All.h>
 #include <Alembic/Abc/ErrorHandler.h>
 #include "ABCTree.h"
+#include "zeno/utils/format.h"
+#include "zeno/utils/string.h"
 #include <numeric>
 #include <filesystem>
 
@@ -211,7 +213,7 @@ struct WriteAlembic2 : INode {
     void write_attrs(std::shared_ptr<PrimitiveObject> prim, T1& schema, T2& samp) {
         OCompoundProperty arbAttrs = schema.getArbGeomParams();
         prim->verts.foreach_attr([&](auto const &key, auto &arr) {
-            if (key == "v" || key == "nrm") {
+            if (key == "v" || key == "nrm" || key == "faceset") {
                 return;
             }
             using T = std::decay_t<decltype(arr[0])>;
@@ -241,6 +243,9 @@ struct WriteAlembic2 : INode {
     void write_user_data(std::shared_ptr<PrimitiveObject> prim, OCompoundProperty& user) {
         auto &ud = prim->userData();
         for (const auto& [key, value] : ud.m_data) {
+            if (key == "faceset_count" || zeno::starts_with(key, "faceset_")) {
+                continue;
+            }
             if (ud.has<int>(key)) {
                 if (user_attrs.count(key) == 0) {
                     auto p = OInt32Property(user, key);
@@ -343,6 +348,28 @@ struct WriteAlembic2 : INode {
         if (prim->polys.size() || prim->tris.size()) {
             // Create a PolyMesh class.
             OPolyMeshSchema &mesh = meshyObj.getSchema();
+            auto &ud = prim->userData();
+            std::vector<std::string> faceSetNames;
+            std::vector<std::vector<int>> faceset_idxs;
+            if (ud.has<int>("faceset_count")) {
+                int faceset_count = ud.get2<int>("faceset_count");
+                for (auto i = 0; i < faceset_count; i++) {
+                    faceSetNames.emplace_back(ud.get2<std::string>(zeno::format("faceset_{:04}", i)));
+                }
+                faceset_idxs.resize(faceset_count);
+                auto &faceset = prim->polys.attr<int>("faceset");
+                for (auto i = 0; i < faceset.size(); i++) {
+                    faceset_idxs[faceset[i]].push_back(i);
+                }
+                for (auto i = 0; i < faceset_count; i++) {
+                    OFaceSet faceset = mesh.createFaceSet(faceSetNames[i]);
+                    OFaceSetSchema facesetSchema = faceset.getSchema ();
+                    OFaceSetSchema::Sample my_face_set_samp ( faceset_idxs[i] );
+                    // faceset is visible, doesn't change.
+                    facesetSchema.set ( my_face_set_samp );
+                    facesetSchema.setFaceExclusivity ( kFaceSetExclusive );
+                }
+            }
 
             OCompoundProperty user = mesh.getUserProperties();
             write_user_data(prim, user);