Skip to content

Commit 1319bbb

Browse files
gngdbgngdb
committed
clean up and set process=False to fix
1 parent af8b08e commit 1319bbb

File tree

1 file changed

+132
-69
lines changed

1 file changed

+132
-69
lines changed

‎transfer_model/write_obj.py

+132-69
Original file line numberDiff line numberDiff line change
@@ -8,78 +8,113 @@
88
import trimesh
99

1010
import smplx
11+
from smplx.joint_names import Body
1112

1213
from tqdm.auto import tqdm, trange
1314

1415
from pathlib import Path
1516

16-
def main(model_folder,
17-
motion_file,
18-
output_folder,
19-
model_type='smplh',
20-
ext='npz',
21-
gender='neutral',
22-
plot_joints=False,
23-
num_betas=10,
24-
sample_expression=True,
25-
num_expression_coeffs=10,
26-
use_face_contour=False):
17+
18+
def main(
19+
model_folder,
20+
motion_file,
21+
output_folder,
22+
model_type="smplh",
23+
ext="npz",
24+
gender="neutral",
25+
plot_joints=False,
26+
num_betas=10,
27+
sample_expression=True,
28+
num_expression_coeffs=10,
29+
use_face_contour=False,
30+
):
2731
output_folder = Path(output_folder)
2832
assert output_folder.exists()
2933

3034
# open motion file
31-
motion = np.load(motion_file)
32-
for k,v in motion.items():
33-
print(k, motion[k].shape)
34-
35-
num_betas = len(motion['betas'])
35+
motion = np.load(motion_file, allow_pickle=True)
36+
for k, v in motion.items():
37+
if type(v) is float:
38+
print(k, v)
39+
else:
40+
print(k, v.shape)
41+
42+
if "betas" in motion:
43+
betas = motion["betas"]
44+
else:
45+
betas = np.zeros((num_betas,))
46+
num_betas = len(betas)
3647
# don't know where this is documented but it's from this part of amass
3748
# https://github.com/nghorbani/amass/blob/master/src/amass/data/prepare_data.py#L39-L40
3849
# gdr2num = {'male':-1, 'neutral':0, 'female':1}
3950
# gdr2num_rev = {v:k for k,v in gdr2num.items()}
40-
gender = str(motion['gender'])
51+
if "gender" in motion:
52+
gender = str(motion["gender"])
53+
else:
54+
gender = gender
55+
4156
print(gender)
4257

4358
print(num_betas)
44-
model = smplx.create(model_folder, model_type=model_type,
45-
gender=gender, use_face_contour=use_face_contour,
46-
num_betas=num_betas,
47-
num_expression_coeffs=num_expression_coeffs,
48-
use_pca=False,
49-
ext=ext)
50-
51-
betas, expression = torch.tensor(motion['betas']).float(), None
52-
betas = betas.unsqueeze(0)[:, :model.num_betas]
53-
poses = torch.tensor(motion['poses']).float()
59+
model = smplx.create(
60+
model_folder,
61+
model_type=model_type,
62+
gender=gender,
63+
use_face_contour=use_face_contour,
64+
num_betas=num_betas,
65+
num_expression_coeffs=num_expression_coeffs,
66+
use_pca=False,
67+
ext=ext,
68+
)
69+
70+
betas, expression = torch.tensor(betas).float(), None
71+
betas = betas.unsqueeze(0)[:, : model.num_betas]
72+
if "poses" in motion:
73+
poses = torch.tensor(motion["poses"]).float()
74+
elif "smpl_poses" in motion:
75+
poses = motion["smpl_poses"]
76+
n = poses.shape[0]
77+
if model_type == "smplh":
78+
poses = np.stack(
79+
[Body.from_smpl(p.reshape(-1, 3)).as_smplh() for p in poses]
80+
)
81+
poses = torch.tensor(poses.reshape(n, -1)).float()
5482
global_orient = poses[:, :3]
55-
body_pose = poses[:, 3:66]
56-
left_hand_pose = poses[:, 66:111]
57-
right_hand_pose = poses[:, 111:]
83+
if model_type == "smplh":
84+
body_pose = poses[:, 3:66]
85+
left_hand_pose = poses[:, 66:111]
86+
right_hand_pose = poses[:, 111:156]
87+
else:
88+
body_pose = poses[:, 3:]
89+
left_hand_pose = np.zeros((n, 3))
90+
right_hand_pose = np.zeros((n, 3))
5891
# if sample_expression:
5992
# expression = torch.randn(
6093
# [1, model.num_expression_coeffs], dtype=torch.float32)
6194

62-
#print(expression)
63-
#print(betas.shape, body_pose.shape, expression.shape)
95+
# print(expression)
96+
# print(betas.shape, body_pose.shape, expression.shape)
6497
for pose_idx in trange(body_pose.size(0)):
6598
pose_idx = [pose_idx]
6699
# output = model(betas=betas, # expression=expression,
67100
# return_verts=True)
68101
output = model(
69-
betas=betas,
70-
global_orient=global_orient[pose_idx],
71-
body_pose=body_pose[pose_idx],
72-
left_hand_pose=left_hand_pose[pose_idx],
73-
right_hand_pose=right_hand_pose[pose_idx],
74-
# expression=expression,
75-
return_verts=True
76-
)
102+
betas=betas,
103+
global_orient=global_orient[pose_idx],
104+
body_pose=body_pose[pose_idx],
105+
left_hand_pose=left_hand_pose[pose_idx],
106+
right_hand_pose=right_hand_pose[pose_idx],
107+
# expression=expression,
108+
return_verts=True,
109+
)
77110
vertices = output.vertices.detach().cpu().numpy().squeeze()
78111
joints = output.joints.detach().cpu().numpy().squeeze()
79112

80113
vertex_colors = np.ones([vertices.shape[0], 4]) * [0.3, 0.3, 0.3, 0.8]
81-
tri_mesh = trimesh.Trimesh(vertices, model.faces,
82-
vertex_colors=vertex_colors)
114+
# process=False to avoid creating a new mesh
115+
tri_mesh = trimesh.Trimesh(
116+
vertices, model.faces, vertex_colors=vertex_colors, process=False
117+
)
83118

84119
output_path = output_folder / "{0:04d}.obj".format(pose_idx[0])
85120
tri_mesh.export(str(output_path))
@@ -102,35 +137,57 @@ def main(model_folder,
102137
pyrender.Viewer(scene, use_raymond_lighting=True)
103138

104139

105-
if __name__ == '__main__':
106-
parser = argparse.ArgumentParser(description='SMPL-X Demo')
107-
108-
parser.add_argument('--model-folder', required=True, type=str,
109-
help='The path to the model folder')
110-
parser.add_argument('--motion-file', required=True, type=str,
111-
help='The path to the motion file to process')
112-
parser.add_argument('--output-folder', required=True, type=str,
113-
help='The path to the output folder')
114-
parser.add_argument('--model-type', default='smplh', type=str,
115-
choices=['smpl', 'smplh', 'smplx', 'mano', 'flame'],
116-
help='The type of model to load')
117-
parser.add_argument('--num-expression-coeffs', default=10, type=int,
118-
dest='num_expression_coeffs',
119-
help='Number of expression coefficients.')
120-
parser.add_argument('--ext', type=str, default='npz',
121-
help='Which extension to use for loading')
122-
parser.add_argument('--sample-expression', default=True,
123-
dest='sample_expression',
124-
type=lambda arg: arg.lower() in ['true', '1'],
125-
help='Sample a random expression')
126-
parser.add_argument('--use-face-contour', default=False,
127-
type=lambda arg: arg.lower() in ['true', '1'],
128-
help='Compute the contour of the face')
140+
if __name__ == "__main__":
141+
parser = argparse.ArgumentParser(description="SMPL-X Demo")
142+
143+
parser.add_argument(
144+
"--model-folder", required=True, type=str, help="The path to the model folder"
145+
)
146+
parser.add_argument(
147+
"--motion-file",
148+
required=True,
149+
type=str,
150+
help="The path to the motion file to process",
151+
)
152+
parser.add_argument(
153+
"--output-folder", required=True, type=str, help="The path to the output folder"
154+
)
155+
parser.add_argument(
156+
"--model-type",
157+
default="smplh",
158+
type=str,
159+
choices=["smpl", "smplh", "smplx", "mano", "flame"],
160+
help="The type of model to load",
161+
)
162+
parser.add_argument(
163+
"--num-expression-coeffs",
164+
default=10,
165+
type=int,
166+
dest="num_expression_coeffs",
167+
help="Number of expression coefficients.",
168+
)
169+
parser.add_argument(
170+
"--ext", type=str, default="npz", help="Which extension to use for loading"
171+
)
172+
parser.add_argument(
173+
"--sample-expression",
174+
default=True,
175+
dest="sample_expression",
176+
type=lambda arg: arg.lower() in ["true", "1"],
177+
help="Sample a random expression",
178+
)
179+
parser.add_argument(
180+
"--use-face-contour",
181+
default=False,
182+
type=lambda arg: arg.lower() in ["true", "1"],
183+
help="Compute the contour of the face",
184+
)
129185

130186
args = parser.parse_args()
131187

132188
def resolve(path):
133189
return osp.expanduser(osp.expandvars(path))
190+
134191
model_folder = resolve(args.model_folder)
135192
motion_file = resolve(args.motion_file)
136193
output_folder = resolve(args.output_folder)
@@ -139,6 +196,12 @@ def resolve(path):
139196
num_expression_coeffs = args.num_expression_coeffs
140197
sample_expression = args.sample_expression
141198

142-
main(model_folder, motion_file, output_folder, model_type, ext=ext,
143-
sample_expression=sample_expression,
144-
use_face_contour=args.use_face_contour)
199+
main(
200+
model_folder,
201+
motion_file,
202+
output_folder,
203+
model_type,
204+
ext=ext,
205+
sample_expression=sample_expression,
206+
use_face_contour=args.use_face_contour,
207+
)

0 commit comments

Comments
 (0)