skinning.h

/*
 * Copyright 2023 Adobe. All rights reserved.
 * This file is licensed to you under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License. You may obtain a copy
 * of the License at http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software distributed under
 * the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR REPRESENTATIONS
 * OF ANY KIND, either express or implied. See the License for the specific language
 * governing permissions and limitations under the License.
 */
#pragma once
 
#include <lagrange/Attribute.h>
#include <lagrange/SurfaceMesh.h>
#include <lagrange/attribute_names.h>
#include <lagrange/utils/assert.h>
#include <lagrange/views.h>
 
#include <Eigen/Core>
#include <numeric>
#include <vector>

 
namespace lagrange::internal {

template <typename Scalar, typename Index>
void skinning_deform(
    SurfaceMesh<Scalar, Index>& mesh,
    const lagrange::Attribute<Scalar>& original_vertices,
    const std::vector<Eigen::Transform<Scalar, 3, Eigen::TransformTraits::Affine>>& transforms,
    const Eigen::Matrix<Scalar, Eigen::Dynamic, Eigen::Dynamic>& weights,
    const Eigen::Matrix<Scalar, Eigen::Dynamic, 1>& weight_complement = {})
{
    la_runtime_assert((size_t)weights.cols() == transforms.size());
    la_runtime_assert(weights.rows() == mesh.get_num_vertices());
    la_runtime_assert(
        (weight_complement.rows() == 0) || (weight_complement.rows() == mesh.get_num_vertices()));
 
    const Index num_vertices = mesh.get_num_vertices();
    const size_t num_handles = transforms.size();
 
    auto original_view = lagrange::matrix_view(original_vertices);
    auto verts = lagrange::vertex_ref(mesh);
    verts.setZero();
 
    for (Index v = 0; v < num_vertices; ++v) {
        const Eigen::Matrix<Scalar, 3, 1> orig = original_view.row(v);
        Scalar weight_sum = 0;
        for (size_t h = 0; h < num_handles; ++h) {
            Scalar weight = weights(v, h);
            if (weight > 0) {
                verts.row(v) += weight * (transforms[h] * orig);
            }
            weight_sum += weight;
        }
        if (weight_complement.rows() > 0 && weight_complement(v) > 0) {
            verts.row(v) += (weight_complement(v) * orig);
            weight_sum += weight_complement(v);
        }
        if (weight_sum > 0) {
            verts.row(v) /= weight_sum;
        }
    }
}

template <typename Scalar, typename Index>
void skinning_deform(
    SurfaceMesh<Scalar, Index>& mesh,
    const Attribute<Scalar>& original_vertices,
    const std::vector<Eigen::Transform<Scalar, 3, Eigen::TransformTraits::Affine>>& transforms)
{
    auto weight_id = mesh.get_attribute_id(AttributeName::weight);
    la_runtime_assert(weight_id != invalid_attribute_id());
 
    auto& weight_attr = mesh.template get_attribute<Scalar>(weight_id);
 
    Eigen::Matrix<Scalar, Eigen::Dynamic, Eigen::Dynamic> weights = matrix_view(weight_attr);
    skinning_deform(mesh, original_vertices, transforms, weights);
}
 
 
template <typename Scalar, typename Index>
struct SkinningExtractNResult
{
    Eigen::Matrix<Scalar, Eigen::Dynamic, Eigen::Dynamic> weights;

    Eigen::Matrix<Index, Eigen::Dynamic, Eigen::Dynamic> indices;
};

template <typename Scalar, typename Index>
SkinningExtractNResult<Scalar, Index> skinning_extract_n(
    const Eigen::Matrix<Scalar, Eigen::Dynamic, Eigen::Dynamic>& weights,
    int n,
    const Eigen::Matrix<Scalar, Eigen::Dynamic, 1>& weight_complement =
        Eigen::Matrix<Scalar, Eigen::Dynamic, 1>())
{
    const Index num_vertices = Index(weights.rows());
    const Index num_handles = Index(weights.cols());
    Index num_implicit_handle = 0;
 
    la_runtime_assert(num_handles > 0);
    if (weight_complement.rows() > 0 &&
        weight_complement.maxCoeff() > std::numeric_limits<Scalar>::epsilon()) {
        la_runtime_assert(weight_complement.rows() == num_vertices);
        // the weights do not sum up to 1, so we assume there is an implicit fixed handle
        num_implicit_handle = 1;
    }
 
    // Max number of handles
    const Index num_handles_max = Index(n) - num_implicit_handle;
    const Index num_handles_used = std::min(num_handles, num_handles_max);
 
 
    SkinningExtractNResult<Scalar, Index> result;
    result.weights.resize(num_vertices, n);
    result.indices.resize(num_vertices, n);
    if (num_handles < Index(n)) {
        // only need to initialize in this case
        result.weights.setZero();
        result.indices.setZero();
    }
 
    if (num_handles <= num_handles_max) {
        // we have more handle slots than handles, so just copy them over
        la_runtime_assert(num_handles == num_handles_used);
        for (Index i = 0; i < num_vertices; ++i) {
            for (Index j = 0; j < num_handles; ++j) {
                result.indices(i, j) = j;
                result.weights(i, j) = weights(i, j);
            }
        }
 
    } else {
        // we only pick some of the handles, find and copy the n most important ones
 
        for (Index i = 0; i < num_vertices; ++i) {
            std::vector<Index> tmp(num_handles);
            std::iota(tmp.begin(), tmp.end(), 0);
            std::partial_sort(
                tmp.begin(),
                tmp.begin() + num_handles_used,
                tmp.end(),
                [&weights, i](const Index a, const Index b) -> bool {
                    return weights(i, a) > weights(i, b);
                });
 
            for (Index j = 0; j < num_handles_used; ++j) {
                result.indices(i, j) = tmp[j];
                result.weights(i, j) = weights(i, tmp[j]);
            }
        }
    }
 
    if (num_implicit_handle > 0) {
        // add the implicit handle
        for (Index i = 0; i < num_vertices; ++i) {
            result.indices(i, num_handles_used) = num_handles_used;
            result.weights(i, num_handles_used) = weight_complement(i);
        }
    }
 
    // normalize if we changed the active handles
    if (num_handles > num_handles_max) {
        for (Index i = 0; i < num_vertices; ++i) {
            Scalar s = result.weights.row(i).sum();
            la_runtime_assert(s > 0);
            result.weights.row(i) /= s;
        }
    }
 
    return result;
}

template <typename Scalar, typename Index, typename WeightScalar = Scalar>
lagrange::AttributeId weights_to_mesh_attribute(
    SurfaceMesh<Scalar, Index>& mesh,
    const Eigen::Matrix<WeightScalar, Eigen::Dynamic, Eigen::Dynamic>& weights)
{
    return mesh.template create_attribute<WeightScalar>(
        AttributeName::weight,
        AttributeElement::Vertex,
        AttributeUsage::Vector,
        weights.cols(),
        {weights.data(), size_t(weights.size())});
}

template <
    typename Scalar,
    typename Index,
    typename WeightScalar = Scalar,
    typename WeightIndex = Index>
std::pair<lagrange::AttributeId, lagrange::AttributeId> weights_to_indexed_mesh_attribute(
    SurfaceMesh<Scalar, Index>& mesh,
    const Eigen::Matrix<WeightScalar, Eigen::Dynamic, Eigen::Dynamic>& weights,
    int n)
{
    auto result = skinning_extract_n<WeightScalar, WeightIndex>(weights, n);
    auto bone_id = mesh.template create_attribute<WeightIndex>(
        AttributeName::indexed_joint,
        AttributeElement::Vertex,
        AttributeUsage::Vector,
        result.indices.cols(),
        {result.indices.data(), size_t(result.indices.size())});
 
    auto weight_id = mesh.template create_attribute<WeightScalar>(
        AttributeName::indexed_weight,
        AttributeElement::Vertex,
        AttributeUsage::Vector,
        result.weights.cols(),
        {result.weights.data(), size_t(result.weights.size())});
    return {bone_id, weight_id};
}
 
 
} // namespace lagrange::internal