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)