momentum/pymomentum/tensor_momentum/tensor_parameter_transform.h at export-D94577295 · facebookresearch/momentum · 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
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

#pragma once

#include <pymomentum/array_utility/default_parameter_set.h>

#include <momentum/character/character.h>
#include <momentum/character/inverse_parameter_transform.h>
#include <momentum/character/parameter_transform.h>

#include <ATen/ATen.h>
#include <pybind11/pybind11.h>

namespace pymomentum {

at::Tensor applyParamTransform(
    const momentum::ParameterTransform* paramTransform,
    at::Tensor modelParams);

at::Tensor applyParamTransform(pybind11::object characters, at::Tensor modelParams);

// Gets the parameter sets specified in the model file as a dictionary
// mapping strings to boolean tensors.
std::unordered_map<std::string, at::Tensor> getParameterSets(
    const momentum::ParameterTransform& parameterTransform);

at::Tensor getScalingParameters(const momentum::ParameterTransform& parameterTransform);

at::Tensor getRigidParameters(const momentum::ParameterTransform& parameterTransform);

at::Tensor getAllParameters(const momentum::ParameterTransform& parameterTransform);

at::Tensor getBlendShapeParameters(const momentum::ParameterTransform& parameterTransform);

at::Tensor getFaceExpressionParameters(const momentum::ParameterTransform& parameterTransform);

at::Tensor getPoseParameters(const momentum::ParameterTransform& parameterTransform);

void addParameterSet(
    momentum::ParameterTransform& parameterTransform,
    const std::string& paramSetName,
    const at::Tensor& paramSet);

at::Tensor getParametersForJoints(
    const momentum::ParameterTransform& parameterTransform,
    const std::vector<size_t>& jointIndices);

at::Tensor findParameters(
    const momentum::ParameterTransform& parameterTransform,
    const std::vector<std::string>& parameterNames,
    bool allowMissing = false);

at::Tensor parameterSetToTensor(
    const momentum::ParameterTransform& parameterTransform,
    const momentum::ParameterSet& paramSet);

at::Tensor applyInverseParamTransform(
    const momentum::InverseParameterTransform* invParamTransform,
    at::Tensor jointParams);

std::unique_ptr<momentum::InverseParameterTransform> createInverseParameterTransform(
    const momentum::ParameterTransform& transform);

momentum::ParameterSet tensorToParameterSet(
    const momentum::ParameterTransform& parameterTransform,
    at::Tensor paramSet,
    DefaultParameterSet defaultParamSet = DefaultParameterSet::NoDefault);

// Change flattened joint parameter motion tensor with shape: [... x
// (n_joint_params x 7)] to shape [... x n_joint_params x 7]. Return the same
// tensor if the input is already unflattened. If `unflattened` is not nullptr,
// assign true only if the function does the shape conversion.
at::Tensor unflattenJointParameters(
    const momentum::Character& character,
    at::Tensor tensor_in,
    bool* unflattened = nullptr);

// Change unflattened joint parameter motion tensor with shape [... x
// n_joint_params x 7)] to shape [... x (n_joint_params x 7)]. Return the same
// tensor if the input is already flattened. If `flattened` is not nullptr,
// assign true only if the function does the shape conversion.
at::Tensor flattenJointParameters(
    const momentum::Character& character,
    at::Tensor tensor_in,
    bool* flattened = nullptr);

at::Tensor modelParametersToBlendShapeCoefficients(
    const momentum::Character& character,
    const at::Tensor& modelParameters);

at::Tensor modelParametersToFaceExpressionCoefficients(
    const momentum::Character& character,
    const at::Tensor& modelParameters);

at::Tensor getParameterTransformTensor(const momentum::ParameterTransform& parameterTransform);

// Map model parameters from one character to another
at::Tensor mapModelParameters(
    const at::Tensor& motion_in,
    const momentum::Character& srcCharacter,
    const momentum::Character& tgtCharacter,
    bool verbose);

// Map model parameters from source names to target character
at::Tensor mapModelParameters_names(
    const at::Tensor& motion_in,
    const std::vector<std::string>& parameterNames_in,
    const momentum::Character& character_remap,
    bool verbose);

// Map joint parameters from one character to another
at::Tensor mapJointParameters(
    at::Tensor srcMotion,
    const momentum::Character& srcCharacter,
    const momentum::Character& tgtCharacter);

// Convert uniform noise to model parameters
at::Tensor uniformRandomToModelParameters(
    const momentum::Character& character,
    at::Tensor unifNoise);

// Apply parameter limits to model parameters
at::Tensor applyModelParameterLimits(
    const momentum::Character& character,
    const at::Tensor& modelParams);

} // namespace pymomentum