aboutsummaryrefslogtreecommitdiff
path: root/scripts/artv2/components/base_components/template.py
blob: 2560ad0ebe6810a25c3383c718af7ec85259d54a (plain) (blame)
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
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
# -*- coding: utf-8 -*-
"""
:author:
    Jeremy Ernst
:description:
    This module contains the class for saving and loading templates, which store components, their settings, their
    positions, and their joint name overrides. These can then be loaded in whole or in part, and can have their
    settings, positions, or overrides loaded onto existing components as well.
"""

import artv2.tools.system.logger.output_logger as logger
import artv2.utilities.general_utilities as utils
import artv2.utilities.component_utilities as component_utils
import artv2.components.base_components.rig_asset as rig_asset
import artv2.components.root as root
import pymel.core as pm


# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
class Template(object):
    """
    The template class saves and loads templates. Templates contain a dictionary of components, their settings, their
    positions, and their joint name overrides.
    """

    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
    def __init__(self, path=None):

        self.logger = logger.OutputLogger(self.__class__.__name__)

        self.path = utils.path_unify(path)
        self.load_settings = True
        self.load_positions = True
        self.load_overrides = True

    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
    def save_template(self):
        """
        Saves a template file to disk.

        example usage:
            .. code-block:: python

                # get a filename to save a template to, then create a template instance and save the template.
                path = QtWidgets.QFileDialog.getSaveFileName(self, "Save Template", template_path, "*.json")
                if os.path.exists(os.path.dirname(path[0])):
                    template_instance = template.Template(utils.path_unify(path[0]))
                    template_instance.save_template()
        """

        self.logger.info("{0}.save_template: Saving template to {1}.".format(self.__class__.__name__, self.path))
        template_data = self._create_template_data()
        utils.save_to_file(self.path, template_data)

    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
    def load_template(self, components, aim=False):
        """
        Loads a template file from disk. If the components don't exist (from the passed in list), then those components
        are created. Then, depending on the settings (load settings, positions, overrides), data is applied to those
        components.

        example usage:
            .. code-block:: python

                # create a template instance using an existing path to a template on disk.
                inst = template.Template(r"C:\\Users\\jernst\\Documents\\ARTv2\\resources\\templates\\test.json")
                inst.load_settings=False
                inst.load_template(["clavicle_l", "arm_l", "torso", "master"])

        :param components: A list of the components from the template to load.
        :param aim: Whether or not aim model should be turned on the component.
        """

        self.logger.info("{0}.load_template: Loading template (from {1}) onto the following components: {2}"
                         "".format(self.__class__.__name__, self.path, components))

        template_data = utils.load_from_file(self.path)

        # ensure an asset and master exist in the scene. If not, create them
        assets = component_utils.get_all_rig_asset_nodes()
        if not assets:
            rig_asset.RigAsset()
            root.Root()

        # load components in scene
        self._load_components(template_data, components)

        # load data onto components
        existing_components = component_utils.get_all_component_names()
        for component in components:
            if component in template_data.keys():
                data = template_data.get(component)
                inst = component_utils.get_component_instance(existing_components.get(component))

                if self.load_settings:
                    self._load_settings(inst, data)

                if self.load_positions:
                    mover_data = data.get("POSITIONS")
                    inst.joint_mover.paste_transforms(template=True, data=mover_data)
                    if aim:
                        inst.joint_mover.aim_helper.toggle_aim_mode()

                if self.load_overrides:
                    self._set_overrides(inst, data)
                del inst

        # set parents
        self._set_parents(components, template_data)
        self._set_mirrors(components, template_data)

    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
    def _set_overrides(self, inst, data):

        self.logger.info("{0}._set_overrides: Setting joint name override data on {1}"
                         "".format(self.__class__.__name__, inst))

        override_data = data.get("OVERRIDES")
        movers = inst.joint_mover.get_movers()
        offset_movers = movers.get("offset")

        for offset in offset_movers:
            if offset.nodeName() in override_data.keys():
                connected_joint = offset.connected_joint.connections(type="transform")[0]
                created_joint = offset.created_joint.get()
                old_name = inst.prefix + connected_joint.otherType.get() + inst.suffix
                if created_joint != override_data.get(offset.nodeName()):
                    inst.rename_joint(override_data.get(offset.nodeName()), old_name, offset)

    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
    def _set_parents(self, components, template_data):

        self.logger.info("{0}._set_parents: Settings parents for {1}".format(self.__class__.__name__, components))

        existing_components = component_utils.get_all_component_names()
        for component in components:
            data = template_data.get(component)
            inst = component_utils.get_component_instance(existing_components.get(component))

            settings_data = data.get("SETTINGS")
            if not settings_data.get("parent") is None:
                existing_joints = component_utils.get_all_created_joints()
                if settings_data.get("parent") in existing_joints:
                    inst.parent = settings_data.get("parent")
                else:
                    inst.parent = "root"

            del inst

    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
    def _set_mirrors(self, components, template_data):

        self.logger.info("{0}._set_mirrors: Setting mirrors for {1}".format(self.__class__.__name__, components))

        existing_components = component_utils.get_all_component_names()

        for component in components:
            data = template_data.get(component)
            inst = component_utils.get_component_instance(existing_components.get(component))

            settings_data = data.get("SETTINGS")
            if settings_data.get("has_mirror"):
                if settings_data.get("mirror_component") is not None:
                    inst.set_mirror(pm.PyNode(settings_data.get("mirror_component")))

            del inst

    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
    def _load_settings(self, instance, data):

        self.logger.info("{0}._load_settings: Loading settings onto {1}".format(self.__class__.__name__, instance))

        settings_data = data.get("SETTINGS")
        properties = utils.get_class_properties(instance)

        if "side" in settings_data.keys():
            if settings_data.get("side") == 1:
                instance.set_side("right")

        for key in settings_data:
            if key in properties:
                setattr(instance, key, settings_data.get(key))

    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
    def _load_components(self, template_data, components):

        existing_components = component_utils.get_all_component_names()
        for component in components:
            if component in template_data.keys():

                data = template_data.get(component)
                info = data.get("INFO")
                settings_data = data.get("SETTINGS")

                # check if component exists in scene. If not, create it
                if component not in existing_components.keys():
                    self.logger.info("{0}._load_components: {1} does not exist. Creating component."
                                     "".format(self.__class__.__name__, component))

                    mod = __import__(info.get("PATH"), {}, {}, [info.get("MODULE")])
                    module_class = getattr(mod, info.get("CLASS"))
                    module_class(prefix=settings_data.get("prefix"), suffix=settings_data.get("suffix"))

    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
    def _create_template_data(self):

        self.logger.info("{0}._create_template_data: Gathering data to build template.".format(self.__class__.__name__))

        template = {}

        # get all component instances
        components = component_utils.return_all_components()

        # for each component, get its settings, mover positions, and joint name overrides
        for component in components:
            inst = component_utils.get_component_instance(component)

            # CLASS INFO
            class_name = component.className.get()
            module_path = component.modulePath.get()
            module_name = module_path.rpartition(".")[2]

            class_info = {"CLASS": class_name, "PATH": module_path, "MODULE": module_name}

            # SETTINGS
            settings = utils.get_class_properties(inst)
            settings["prefix"] = inst.prefix
            settings["suffix"] = inst.suffix
            settings["parent"] = inst.parent
            settings["has_mirror"] = component.has_mirror.get()
            settings["mirror_component"] = None
            if component.has_mirror.get():
                settings["mirror_component"] = component.mirror_component.connections()[0].nodeName()
            if inst.has_sides:
                settings["side"] = component.side.get()

            # MOVER POSITIONS
            mover_positions = inst.joint_mover.copy_transforms(template=True)

            # JOINT NAME OVERRIDES
            movers = inst.joint_mover.get_movers()
            offset_movers = movers.get("offset")
            override_data = {}
            for mover in offset_movers:

                driven_joint = mover.connected_joint.connections(type="transform")[0]
                original_name = inst.prefix + driven_joint.otherType.get() + inst.suffix
                if mover.created_joint.get() != original_name:
                    override_data[mover.nodeName()] = mover.created_joint.get()

            template[component.componentName.get()] = {"SETTINGS": settings,
                                                       "INFO": class_info,
                                                       "POSITIONS": mover_positions,
                                                       "OVERRIDES": override_data}

            del inst
        return template