自己动手实现SMPL算法

2024-06-15

字数统计: 999字 | 阅读时长: 5分

最近在学习3D相关的算法，遇到SMPL算法，不甚理解，所以打算自己实现一遍，毕竟————学习新知识的最好方法就是自己实现一下。
项目地址：https://github.com/killf/smpl

一、代码

model.py是加载模型文件和算法的入口。

import numpy as np
import pickle

from .lbs import blend_skin
from .utils import rodrigues, batch_eye


class SMPL:
    def __init__(self, model_file):
        model = pickle.load(open(model_file, "rb"), encoding='iso-8859-1')
        self.v_template = model["v_template"]  # 平均body
        self.weights = model["weights"]  # joints对顶点的权重
        self.posedirs = model["posedirs"]  # pose对顶点的影响
        self.shapedirs = model["shapedirs"].x  # shape对顶点的影响
        self.J_regressor = model["J_regressor"]  # joints回归器，根据顶点坐标估算joints坐标
        self.faces = model["f"] + 1  # faces的索引，索引从1开始

        kintree_table = model["kintree_table"]  # 运动学树
        id_to_col = {kintree_table[1, i]: i for i in range(kintree_table.shape[1])}
        self.kintree_table = {i: id_to_col[kintree_table[0, i]] for i in range(1, kintree_table.shape[1])}  # 关节点之间的父子级关系，0->父，1->子

    def __call__(self, shape=None, pose=None, trans=None):
        """
        SMPL算法
        :param shape: 10
        :param pose: 24x3
        :param trans: 3
        """
        shape = np.zeros(10) if shape is None else shape
        pose = np.zeros(24, 3) if pose is None else pose
        trans = np.zeros(3) if trans is None else trans

        v_shaped = self.v_template + self.shapedirs.dot(shape)  # shape->顶点
        v_posed = v_shaped + self.posedirs.dot((rodrigues(pose[1:]) - batch_eye(3)).ravel())  # pose->顶点

        joints = self.J_regressor.dot(v_shaped)  # 根据顶点估算的joints坐标
        vertices = blend_skin(pose, v_posed, joints, self.weights, self.kintree_table)  # 姿态影响的顶点变化
        vertices = vertices + trans.reshape((1, 3))

        joints = self.J_regressor.dot(vertices)  # 计算 pose 下 joints 位置
        return vertices, joints

lbs.py是对lbs算法的实现，即线性混合蒙皮算法。

import numpy as np

from .utils import *


def rigid_transformation(pose, joints, kintree_table):
    # 首先计算根结点 (0) 的世界坐标变换, 或者说是根结点相对世界坐标系的位姿
    T = np.zeros([4, 4])
    T[:3, :3] = rodrigues(pose[0])  # 轴角转换到旋转矩阵，相对世界坐标
    T[:3, 3] = joints[0]  # 根结点在世界坐标系下的位置
    T[3, 3] = 1  # 齐次矩阵，1

    Ts = np.zeros([24, 4, 4])
    Ts[0] = T

    # 计算子节点 (1~24) 的世界坐标变换
    for i in range(1, 24):
        # 首先计算子节点相对父节点坐标系的位姿 [R|t]
        T = np.zeros([4, 4])
        T[3, 3] = 1
        # 计算子节点相对父节点的旋转矩阵 R
        T[:3, :3] = rodrigues(pose[i])
        # 计算子节点相对父节点的偏移量 t
        T[:3, 3] = joints[i] - joints[kintree_table[i]]
        # 然后计算子节点相对世界坐标系的位姿
        Ts[i] = np.matmul(Ts[kintree_table[i]], T)  # 乘上其父节点的变换矩阵

    # 计算每个子节点相对 T-pose 时的位姿矩阵
    for i in range(24):
        R = Ts[i][:3, :3]
        t = Ts[i][:3, 3] - R.dot(joints[i])  # 子节点相对T-pose的偏移
        Ts[i][:3, 3] = t

    return Ts


def blend_skin(pose, v_posed, joints, weights, kintree_table):
    # 计算每个关节点相对于默认姿态的旋转向量和平移向量
    Ts = rigid_transformation(pose, joints, kintree_table)

    # 开始蒙皮操作，LBS 过程
    v_posed_homo = np.hstack([v_posed, np.ones([v_posed.shape[0], 1])])
    vertices_homo = np.matmul(weights.dot(Ts.reshape([24, 16])).reshape([-1, 4, 4]), v_posed_homo.reshape([-1, 4, 1]))
    vertices = vertices_homo[:, :3, 0]  # 由于是齐次矩阵，取前3列
    return vertices

utils.py是一些工具代码。

import numpy as np


def rodrigues_np(r):
    theta = np.linalg.norm(r)
    x, y, z = r / theta
    c, s = np.cos(theta), np.sin(theta)
    return np.array([[x * x * (1 - c) + c, x * y * (1 - c) + z * s, x * z * (1 - c) - y * s],
                     [x * y * (1 - c) - z * s, y * y * (1 - c) + c, y * z * (1 - c) + x * s],
                     [x * z * (1 - c) + y * s, y * z * (1 - c) - x * s, z * z * (1 - c) + c]]).T


def rodrigues(x):
    if x.ndim == 1:
        return rodrigues_np(x)
    else:
        return np.stack([rodrigues(i) for i in x])


def batch_eye(n):
    return np.eye(n)[np.newaxis, ...]


def write_obj(file, vertices, faces):
    with open(file, 'w') as fp:
        for v in vertices:
            fp.write('v %f %f %f\n' % (v[0], v[1], v[2]))
        for f in faces:
            fp.write('f %d %d %d\n' % (f[0], f[1], f[2]))

二、测试

import numpy as np
from smpl import SMPL, write_obj

model = SMPL("../models/basicmodel_m_lbs_10_207_0_v1.0.0.pkl")

pose = (np.random.rand(24, 3) - 0.5) * .4
betas = (np.random.rand(10) - 0.5) * .06
vertices, _ = model(betas, pose)
write_obj("test.obj", vertices, model.faces)

body

本文作者： killf
本文链接： http://www.killf.info/机器学习/自己动手实现SMPL算法/
版权声明： 本博客所有文章除特别声明外，均采用 Apache License 2.0 许可协议。转载请注明出处！