body-pose-animation/model.py
2021-02-01 23:53:36 +01:00

92 lines
2.6 KiB
Python

import matplotlib.pyplot as plt
import numpy as np
import smplx
from human_body_prior.body_model.body_model_vposer import BodyModelWithPoser
class VPoserModel():
def __init__(
self,
model_type='smpl',
vposer_model_path="./vposer_v1_0",
ext='npz',
gender='neutral',
create_body_pose=True,
plot_joints=True,
num_betas=10,
sample_shape=False,
sample_expression=False,
num_expression_coeffs=10,
use_face_contour=False
):
self.vposer_model_path = vposer_model_path
self.model_type = model_type
self.ext = ext
self.gender = gender
self.plot_joints = plot_joints
self.num_betas = num_betas
self.sample_shape = sample_shape
self.sample_expression = sample_expression
self.num_expression_coeffs = num_expression_coeffs
self.create_body_pose = create_body_pose
self.create_model()
def create_model(self):
self.model = BodyModelWithPoser(
bm_path="./models/smplx/SMPLX_MALE.npz",
batch_size=1,
poser_type="vposer",
smpl_exp_dir=self.vposer_model_path
)
return self.model
def get_vposer_latens(self):
return self.model.poZ_body
def get_pose(self):
return self.model.pose_body
class SMPLyModel():
def __init__(
self,
model_folder,
model_type='smplx',
ext='npz',
gender='male',
create_body_pose=True,
plot_joints=True,
num_betas=10,
sample_shape=False,
sample_expression=True,
num_expression_coeffs=10,
plotting_module='pyrender',
use_face_contour=False
):
self.model_folder = model_folder
self.model_type = model_type
self.ext = ext
self.gender = gender
self.plot_joints = plot_joints
self.num_betas = num_betas
self.sample_shape = sample_shape
self.sample_expression = sample_expression
self.num_expression_coeffs = num_expression_coeffs
self.plotting_module = plotting_module
self.use_face_contour = use_face_contour
self.create_body_pose = create_body_pose
def create_model(self):
self.model = smplx.create(
self.model_folder,
model_type=self.model_type,
gender=self.gender,
use_face_contour=self.use_face_contour,
create_body_pose=self.create_body_pose,
num_betas=self.num_betas,
return_verts=True,
num_expression_coeffs=self.num_expression_coeffs,
ext=self.ext)
return self.model