Spaces:
Running
Running
File size: 8,769 Bytes
15d6c34 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 |
import argparse
import logging
import logging as log
import os
import time
from collections import defaultdict
from os.path import join as pjoin
from typing import Dict, Optional, Tuple
import imageio
import numpy as np
import pyrender
import smplx
import torch
import trimesh
from numpy.typing import ArrayLike
from torch import Tensor
from tqdm import tqdm
from .motionx_explorer import (NUM_FACIAL_EXPRESSION_DIMS,
calc_mean_stddev_pose, get_info_from_file,
label_code, motion_arr_to_dict, names_to_arrays,
to_smplx_dict)
log.basicConfig(
level=log.INFO,
format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
)
def save_img(img, save_path):
imageio.imwrite(save_path, img)
# based on https://github.com/vchoutas/smplx/blob/main/examples/demo.py
# used to render one pose (not sequence of poses) e.g. to see the mean pose
def render_mesh(model, output, should_save=False, save_path=None):
should_display = not should_save
vertices = output.vertices.detach().cpu().numpy().squeeze()
# joint points not visualized for now
# joints = output.joints.detach().cpu().numpy().squeeze()
scene = pyrender.Scene()
if should_display:
viewer = pyrender.Viewer(scene, run_in_thread=True)
mesh_node = None
joints_node = None
# Rotation matrix (90 degrees around the X-axis)
rot = trimesh.transformations.rotation_matrix(np.radians(90), [1, 0, 0])
if should_save:
os.makedirs(os.path.dirname(save_path), exist_ok=True)
try:
# print("Vertices shape =", vertices.shape)
# print("Joints shape =", joints.shape)
# from their demo script
plotting_module = "pyrender"
if plotting_module == "pyrender":
vertex_colors = np.ones([vertices.shape[0], 4]) * [0.3, 0.3, 0.3, 0.8]
tri_mesh = trimesh.Trimesh(vertices, model.faces, vertex_colors=vertex_colors)
# Apply rotation
tri_mesh.apply_transform(rot)
##### RENDER LOCK #####
if should_display:
viewer.render_lock.acquire()
if mesh_node:
scene.remove_node(mesh_node)
mesh = pyrender.Mesh.from_trimesh(tri_mesh)
mesh_node = scene.add(mesh)
camera = pyrender.PerspectiveCamera(yfov=np.pi / 3.0, aspectRatio=1.0)
min_bound, max_bound = mesh.bounds
# Calculate the center of the bounding box
center = (min_bound + max_bound) / 2
# Calculate the extents (the dimensions of the bounding box)
extents = max_bound - min_bound
# Estimate a suitable distance
distance = max(extents) * 2 # Adjust the multiplier as needed
# Create a camera pose matrix
cam_pose = np.array(
[
[1.0, 0, 0, center[0]],
[0, 1.0, 0, center[1]-1.0],
[0, 0, 1.0, center[2] + distance + 0.5],
[0, 0, 0, 1],
]
)
# Rotate around X-axis
angle = np.radians(90)
cos_angle = np.cos(angle)
sin_angle = np.sin(angle)
rot_x = np.array([
[1, 0, 0, 0],
[0, cos_angle, -sin_angle, 0],
[0, sin_angle, cos_angle, 0],
[0, 0, 0, 1]
])
cam_pose = np.matmul(cam_pose, rot_x)
# this is great pose, head on, but a bit far from face
# cam_pose[:3, 3] += np.array([0, 0, -3.5])
cam_pose[:3, 3] += np.array([-.01, 0.65, -3.3])
scene.add(camera, pose=cam_pose)
# Add light for better visualization
light = pyrender.DirectionalLight(color=np.ones(3), intensity=2.0)
scene.add(light, pose=cam_pose)
if should_save:
r = pyrender.OffscreenRenderer(viewport_width=640, viewport_height=480)
col_img, _ = r.render(scene)
save_img(col_img, save_path)
r.delete() # Free up the resources
###### RENDER LOCK RELEASE #####
if should_display:
viewer.render_lock.release()
except KeyboardInterrupt:
if should_display:
viewer.close_external()
# motion_arr is 212 dims (no shapes: aka no betas and no face shapes)
def mesh_and_save(args, motion_arr, seq_name, model_name, emotion):
motion_dict = motion_arr_to_dict(motion_arr, shapes_dropped=True)
smplx_params = to_smplx_dict(motion_dict)
model_folder = "./models/smplx"
batch_size = 1
model = smplx.SMPLX(
model_folder,
use_pca=False, # our joints are not in pca space
num_expression_coeffs=NUM_FACIAL_EXPRESSION_DIMS,
batch_size=batch_size,
)
output = model.forward(**smplx_params, return_verts=True)
log.info(f"output size {output.vertices.shape}")
log.info(f"output size {output.joints.shape}")
log.info("rendering mesh")
base_file = args.file.split('.')[0]
# add {emotion}_{base_file} as a subfolder if it doesn't exist
subfolder = f"single_pose_imgs/{model_name}/{emotion}_{base_file}"
if not os.path.exists(subfolder):
os.makedirs(subfolder)
save_path = f"{subfolder}/{seq_name}_pose.png"
render_mesh(model, output, should_save=True, save_path=save_path)
log.warning(
"if you don't see the mesh animation, make sure you are running on graphics compatible DTU machine (vgl xterm)."
)
return subfolder
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument(
"-e",
"--emotion",
type=str,
required=True,
help="emotion to calculate mean, std for",
)
parser.add_argument(
"-f",
"--file",
type=str,
required=True,
help="file to filter for emotion",
)
parser.add_argument(
"-m",
"--model_path",
type=str,
required=False,
default="",
help="Path to model directory e.g. ./checkpoints/grab/grab_baseline_dp_2gpu_8layers_1000",
)
args = parser.parse_args()
data_root = './data/GRAB'
motion_label_dir = pjoin(data_root, 'texts')
emotions_label_dir = pjoin(data_root, 'face_texts')
args = parser.parse_args()
seq_list_file = pjoin(data_root, args.file)
logging.info("aggregating info about sequences...")
info_dict = get_info_from_file(seq_list_file, emotions_label_dir, motion_label_dir)
# get all files with args.emotion_code
logging.info("calculating mean pose statistics...")
emotions = info_dict["unique_emotions"]
# emotions = [args.emotion]
for emotion in emotions:
logging.info(f"render mean mesh for {emotion} in {args.file}...")
emo_code = label_code(emotion)
names_with_emo = info_dict["emotion_to_names"][emo_code]
arrays = names_to_arrays(data_root, names_with_emo)
mean, std = calc_mean_stddev_pose(arrays)
# add 1 dimension to mean and std
mean = mean.reshape(1, -1)
std = std.reshape(1, -1)
mean_dict = motion_arr_to_dict(mean, shapes_dropped=True)
std_dict = motion_arr_to_dict(std, shapes_dropped=True)
logging.info(f"{emotion} mean: {mean_dict['face_expr']}")
logging.info(f"{emotion} std: {std_dict['face_expr']}")
logging.info(f"rendering mean mesh for {emotion} in {args.file}...")
subfolder = mesh_and_save(args, mean, "mean", args.model_path, emotion)
model_name = args.model_path.split('/')[-1] if args.model_path else "ground_truth"
# write the sequence names in a metadata folder at subfolder
metadata_folder = f"{subfolder}/metadata"
if not os.path.exists(metadata_folder):
os.makedirs(metadata_folder)
metadata_path = f"{metadata_folder}/metadata.txt"
with open(metadata_path, 'w') as f:
f.write(f"model: {model_name}\n")
f.write(f"emotion: {emotion}\n")
f.write(f"file: {args.file}\n")
f.write(f"mean: {mean_dict}\n")
f.write(f"std: {std_dict}\n")
for name in names_with_emo:
f.write(f"{name}\n")
# now plot mesh for each of the sequences
for i, arr in enumerate(arrays):
one_pose = arr[0]
one_pose = one_pose.reshape(1, -1)
name = names_with_emo[i]
# replace / with _
name = name.replace("/", "_")
subfolder = mesh_and_save(args, one_pose, name, args.model_path, emotion)
|