|
| 1 | +/* |
| 2 | + * Copyright (c) Meta Platforms, Inc. and affiliates. |
| 3 | + * |
| 4 | + * This source code is licensed under the MIT license found in the |
| 5 | + * LICENSE file in the root directory of this source tree. |
| 6 | + */ |
| 7 | + |
| 8 | +#include "pymomentum/geometry/parameter_transform_pybind.h" |
| 9 | + |
| 10 | +#include "pymomentum/geometry/momentum_geometry.h" |
| 11 | +#include "pymomentum/tensor_momentum/tensor_parameter_transform.h" |
| 12 | + |
| 13 | +#include <momentum/character/inverse_parameter_transform.h> |
| 14 | +#include <momentum/character/parameter_transform.h> |
| 15 | + |
| 16 | +#include <pybind11/eigen.h> |
| 17 | +#include <pybind11/pybind11.h> |
| 18 | +#include <pybind11/stl.h> |
| 19 | +#include <torch/csrc/utils/pybind.h> |
| 20 | +#include <torch/python.h> |
| 21 | + |
| 22 | +#include <fmt/format.h> |
| 23 | + |
| 24 | +#include <utility> |
| 25 | + |
| 26 | +namespace py = pybind11; |
| 27 | +namespace mm = momentum; |
| 28 | + |
| 29 | +namespace pymomentum { |
| 30 | + |
| 31 | +void registerParameterTransformBindings( |
| 32 | + py::class_<mm::ParameterTransform>& parameterTransformClass) { |
| 33 | + // ===================================================== |
| 34 | + // momentum::ParameterTransform |
| 35 | + // - names |
| 36 | + // - size() |
| 37 | + // - apply(modelParameters) |
| 38 | + // - getScalingParameters() |
| 39 | + // - getRigidParameters() |
| 40 | + // - getParametersForJoints(jointIndices) |
| 41 | + // - createInverseParameterTransform() |
| 42 | + // ===================================================== |
| 43 | + parameterTransformClass |
| 44 | + .def( |
| 45 | + py::init([](const std::vector<std::string>& names, |
| 46 | + const mm::Skeleton& skeleton, |
| 47 | + const Eigen::SparseMatrix<float, Eigen::RowMajor>& transform) { |
| 48 | + mm::ParameterTransform parameterTransform; |
| 49 | + parameterTransform.name = names; |
| 50 | + parameterTransform.transform.resize( |
| 51 | + static_cast<int>(skeleton.joints.size()) * mm::kParametersPerJoint, |
| 52 | + static_cast<int>(names.size())); |
| 53 | + |
| 54 | + for (int i = 0; i < transform.outerSize(); ++i) { |
| 55 | + for (Eigen::SparseMatrix<float, Eigen::RowMajor>::InnerIterator it(transform, i); it; |
| 56 | + ++it) { |
| 57 | + parameterTransform.transform.coeffRef( |
| 58 | + static_cast<long>(it.row()), static_cast<long>(it.col())) = it.value(); |
| 59 | + } |
| 60 | + } |
| 61 | + |
| 62 | + parameterTransform.offsets.setZero(skeleton.joints.size() * mm::kParametersPerJoint); |
| 63 | + return parameterTransform; |
| 64 | + }), |
| 65 | + py::arg("names"), |
| 66 | + py::arg("skeleton"), |
| 67 | + py::arg("transform")) |
| 68 | + .def_readonly("names", &mm::ParameterTransform::name, "List of model parameter names") |
| 69 | + .def_property_readonly( |
| 70 | + "size", |
| 71 | + &mm::ParameterTransform::numAllModelParameters, |
| 72 | + "Size of the model parameter vector.") |
| 73 | + .def( |
| 74 | + "apply", |
| 75 | + [](const mm::ParameterTransform* paramTransform, |
| 76 | + torch::Tensor modelParams) -> torch::Tensor { |
| 77 | + return applyParamTransform(paramTransform, std::move(modelParams)); |
| 78 | + }, |
| 79 | + R"(Apply the parameter transform to a k-dimensional model parameter vector (returns the 7*nJoints joint parameter vector). |
| 80 | +
|
| 81 | +The modelParameters store the reduced set of parameters (typically around 50) that are actually |
| 82 | +optimized in the IK step. |
| 83 | +
|
| 84 | +The jointParameters are stored (tx, ty, tz; rx, ry, rz; s) and each represents the transform relative to the parent joint. |
| 85 | +Rotations are in Euler angles.)", |
| 86 | + py::arg("model_parameters")) |
| 87 | + .def_property_readonly( |
| 88 | + "scaling_parameters", |
| 89 | + &getScalingParameters, |
| 90 | + "Boolean torch.Tensor indicating which parameters are used to control the character's scale.") |
| 91 | + .def_property_readonly( |
| 92 | + "rigid_parameters", |
| 93 | + &getRigidParameters, |
| 94 | + "Boolean torch.Tensor indicating which parameters are used to control the character's rigid transform (translation and rotation).") |
| 95 | + .def_property_readonly( |
| 96 | + "all_parameters", &getAllParameters, "Boolean torch.Tensor with all parameters enabled.") |
| 97 | + .def_property_readonly( |
| 98 | + "blend_shape_parameters", |
| 99 | + &getBlendShapeParameters, |
| 100 | + "Boolean torch.Tensor with just the blend shape parameters enabled.") |
| 101 | + .def_property_readonly( |
| 102 | + "pose_parameters", |
| 103 | + &getPoseParameters, |
| 104 | + "Boolean torch.Tensor with all the parameters used to pose the body, excluding and scaling, blend shape, or physics parameters.") |
| 105 | + .def_property_readonly( |
| 106 | + "no_parameters", |
| 107 | + [](const momentum::ParameterTransform& parameterTransform) { |
| 108 | + return parameterSetToTensor(parameterTransform, momentum::ParameterSet()); |
| 109 | + }, |
| 110 | + "Boolean torch.Tensor with no parameters enabled.") |
| 111 | + .def_property_readonly( |
| 112 | + "parameter_sets", |
| 113 | + &getParameterSets, |
| 114 | + R"(A dictionary mapping names to sets of parameters (as a boolean torch.Tensor) that are defined in the .model file. |
| 115 | +This is convenient for turning off certain body features; for example the 'fingers' parameters |
| 116 | +can be used to enable/disable finger motion in the character model. )") |
| 117 | + .def( |
| 118 | + "parameters_for_joints", |
| 119 | + &getParametersForJoints, |
| 120 | + R"(Gets a boolean torch.Tensor indicating which parameters affect the passed-in joints. |
| 121 | +
|
| 122 | +:param jointIndices: List of integers of skeleton joints.)", |
| 123 | + py::arg("joint_indices")) |
| 124 | + .def( |
| 125 | + "find_parameters", |
| 126 | + &findParameters, |
| 127 | + R"(Return a boolean tensor with the named parameters set to true. |
| 128 | +
|
| 129 | +:param parameter_names: Names of the parameters to find. |
| 130 | +:param allow_missng: If false, missing parameters will throw an exception. |
| 131 | + )", |
| 132 | + py::arg("names"), |
| 133 | + py::arg("allow_missing") = false) |
| 134 | + .def( |
| 135 | + "inverse", |
| 136 | + &createInverseParameterTransform, |
| 137 | + R"(Compute the inverse of the parameter transform (a mapping from joint parameters to model parameters). |
| 138 | +
|
| 139 | +:return: The inverse parameter transform.)") |
| 140 | + .def_property_readonly( |
| 141 | + "transform", |
| 142 | + &getParameterTransformTensor, |
| 143 | + "Returns the parameter transform matrix which when applied maps model parameters to joint parameters.") |
| 144 | + .def("__repr__", [](const mm::ParameterTransform& pt) { |
| 145 | + return fmt::format( |
| 146 | + "ParameterTransform(parameters={}, joints={})", |
| 147 | + pt.numAllModelParameters(), |
| 148 | + pt.transform.rows() / mm::kParametersPerJoint); |
| 149 | + }); |
| 150 | +} |
| 151 | + |
| 152 | +void registerInverseParameterTransformBindings( |
| 153 | + py::class_<mm::InverseParameterTransform>& inverseParameterTransformClass) { |
| 154 | + // ===================================================== |
| 155 | + // momentum::InverseParameterTransform |
| 156 | + // - apply() |
| 157 | + // ===================================================== |
| 158 | + inverseParameterTransformClass |
| 159 | + .def( |
| 160 | + "apply", |
| 161 | + &applyInverseParamTransform, |
| 162 | + R"(Apply the inverse parameter transform to a 7*nJoints-dimensional joint parameter vector (returns the k-dimensional model parameter vector). |
| 163 | +
|
| 164 | +Because the number of joint parameters is much larger than the number of model parameters, this will in general have a non-zero residual. |
| 165 | +
|
| 166 | +:param joint_parameters: Joint parameter tensor with dimensions (nBatch x 7*nJoints). |
| 167 | +:return: A torch.Tensor containing the (nBatch x nModelParameters) model parameters.)", |
| 168 | + py::arg("joint_parameters")) |
| 169 | + .def("__repr__", [](const mm::InverseParameterTransform& ipt) { |
| 170 | + return fmt::format( |
| 171 | + "InverseParameterTransform(parameters={}, joints={})", |
| 172 | + ipt.transform.cols(), |
| 173 | + ipt.transform.rows() / mm::kParametersPerJoint); |
| 174 | + }); |
| 175 | +} |
| 176 | + |
| 177 | +} // namespace pymomentum |
![meta-codesync[bot]](https://avatars.githubusercontent.com/in/1373969?v=4&size=40){kind=avatar}
0 commit comments