UniversalSceneDescriptionTools/multi_variant_usd_tool.py at master · JJGIV2010/UniversalSceneDescriptionTools · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
from pxr import Usd, UsdGeom, UsdPhysics, UsdShade, Sdf, Gf, Tf

class USD_VariantTool:
    def __init__(self):
        self.main_directory_or_url = ""
        self.stage = omni.usd.get_context().get_stage()
        self.variant_dict = {
            "name":"",
            "variant_sets":[], # list of strings
            "variants":[], # list of 2 value tuples (variant_set,variant)
            "variant_custom_attributes":[] # (variant_set,variant,attribute,attribute_value)
        }
        self.sdf_dict = {
            int: Sdf.ValueTypeNames.Int,
            float: Sdf.ValueTypeNames.Float,
            str: Sdf.ValueTypeNames.String
        }

    def add_comment(self,prim_path,text):
        prim = self.stage.GetPrimAtPath(prim_path)
        prim.SetMetadata('comment',text)

    def set_kind(self,path,kind):
        omni.kit.commands.execute('ChangeMetadata',
                                  object_paths=[Sdf.Path(path)],
                                  key='kind',
                                  value=kind)

    def add_variant_set(self,variant_set):
        self.variant_dict['variant_sets'].append(variant_set)

    def add_variant(self,variant_set,variant_name):
        self.variant_dict['variants'].append((variant_set,variant_name))

    def add_custom_attribute(self,variant_set,variant,attribute_name,attribute_value):
        self.variant_dict['variant_custom_attributes'].append((variant_set,variant,attribute_name,attribute_value))

    def create_new_nested_object(self,name,variant_sets,variants,var_attributes):
        self.variant_dict['name'] = name
        for v_set in range(len(variant_sets)):
            self.add_variant_set(variant_sets[v_set])
        for i in range(len(variants)):
            self.add_variant(variants[i][0],variants[i][1])
        for i in range(len(var_attributes)):
            self.add_custom_attribute(var_attributes[i][0],var_attributes[i][1],var_attributes[i][2],var_attributes[i][3])

    def build_new_nested_object(self):
        self.main_xform = UsdGeom.Xform.Define(self.stage, f"/{self.variant_dict['name']}")
        self.main_prim = self.stage.DefinePrim(f"/{self.variant_dict['name']}/{self.variant_dict['name']}_Variants")

        self.set_kind(self.main_xform.GetPath().pathString,"assembly")
        self.add_comment(self.main_xform.GetPath().pathString,
                         f"Main {self.variant_dict['name']} assembly, with {len(self.variant_dict['variant_sets'])} variant sets ")

        self.set_kind(self.main_prim.GetPath().pathString, "group")
        self.add_comment(self.main_prim.GetPath().pathString,
                         f"{self.variant_dict['name']} group, with {len(self.variant_dict['variants'])} components as variants ")

        for i in range(len(object_builder.variant_dict['variant_sets'])):
            for k in range(len(object_builder.variant_dict['variants'])):
                variant_set = self.main_prim.GetVariantSets().AddVariantSet(self.variant_dict['variant_sets'][i])

                if self.variant_dict['variants'][k][0] == self.variant_dict['variant_sets'][i]:
                    variant_set.AddVariant(self.variant_dict['variants'][k][1])
                    variant_set.SetVariantSelection(self.variant_dict["variants"][k][1])

                    with variant_set.GetVariantEditContext():
                        new_variant_prim = self.stage.DefinePrim(self.main_prim.GetPath().AppendChild(
                            f"{self.variant_dict['variants'][k][1]}_mesh"))
                        self.set_kind(new_variant_prim.GetPath().pathString, "component")
                        self.add_comment(new_variant_prim.GetPath().pathString,
                                         f"{self.variant_dict['variants'][k][1]} component custom attributes ")

                        new_variant_prim.GetReferences().AddReference(
                           f"{self.main_directory_or_url}/{self.variant_dict['name']}/{self.variant_dict['variant_sets'][i]}/{self.variant_dict['variants'][i][0]}.usda")
                        for a in range(len(self.variant_dict['variant_custom_attributes'])):
                            if self.variant_dict['variant_custom_attributes'][a][0] == self.variant_dict['variant_sets'][i]:
                                if self.variant_dict['variant_custom_attributes'][a][1] == self.variant_dict['variants'][k][1]:
                                    new_attribute = new_variant_prim.CreateAttribute(
                                        self.variant_dict['variant_custom_attributes'][a][2],
                                        self.sdf_dict[type(self.variant_dict['variant_custom_attributes'][a][3])])
                                    new_attribute.Set(self.variant_dict['variant_custom_attributes'][a][3])

                else:
                    pass

    def _show_as_string(self):
        print(self.stage.GetRootLayer().ExportToString())

object_builder = USD_VariantTool()

object_builder.create_new_nested_object(
    "food", # string
    ["burritos","tacos"], # list of strings
    [
        ("burritos","chicken"), # list of two value tuples
        ("burritos","fish"),
        ("burritos","carnitas"),
        ("tacos", "steak"),
        ("tacos", "chicken")
    ],[
        ("burritos", "chicken","test_int",10), # list of four value tuples
        ("burritos", "fish","test_float",2.0),
        ("burritos", "carnitas","sauce","light"),
        ("tacos", "steak","sauce","none"),
        ("tacos", "chicken","sauce","super spicy")
    ])

object_builder.build_new_nested_object()