import os import numpy as np import random import torch import shutil import csv import pprint import pandas as pd from loguru import logger from collections import OrderedDict import matplotlib.pyplot as plt import pickle import time import hashlib from scipy.spatial.transform import Rotation as R from scipy.spatial.transform import Slerp import cv2 import utils.media import utils.fast_render def write_wav_names_to_csv(folder_path, csv_path): """ Traverse a folder and write the base names of all .wav files to a CSV file. :param folder_path: Path to the folder to traverse. :param csv_path: Path to the CSV file to write. """ # Open the CSV file for writing with open(csv_path, mode='w', newline='') as file: writer = csv.writer(file) # Write the header writer.writerow(['id', 'type']) # Walk through the folder for root, dirs, files in os.walk(folder_path): for file in files: # Check if the file ends with .wav if file.endswith('.wav'): # Extract the base name without the extension base_name = os.path.splitext(file)[0] # Write the base name and type to the CSV writer.writerow([base_name, 'test']) def resize_motion_sequence_tensor(sequence, target_frames): """ Resize a batch of 8-frame motion sequences to a specified number of frames using interpolation. :param sequence: A (bs, 8, 165) tensor representing a batch of 8-frame motion sequences :param target_frames: An integer representing the desired number of frames in the output sequences :return: A (bs, target_frames, 165) tensor representing the resized motion sequences """ bs, _, _ = sequence.shape # Create a time vector for the original and target sequences original_time = torch.linspace(0, 1, 8, device=sequence.device).view(1, -1, 1) target_time = torch.linspace(0, 1, target_frames, device=sequence.device).view(1, -1, 1) # Permute the dimensions to (bs, 165, 8) for interpolation sequence = sequence.permute(0, 2, 1) # Interpolate each joint's motion to the target number of frames resized_sequence = torch.nn.functional.interpolate(sequence, size=target_frames, mode='linear', align_corners=True) # Permute the dimensions back to (bs, target_frames, 165) resized_sequence = resized_sequence.permute(0, 2, 1) return resized_sequence def adjust_speed_according_to_ratio_tensor(chunks): """ Adjust the playback speed within a batch of 32-frame chunks according to random intervals. :param chunks: A (bs, 32, 165) tensor representing a batch of motion chunks :return: A (bs, 32, 165) tensor representing the motion chunks after speed adjustment """ bs, _, _ = chunks.shape # Step 1: Divide the chunk into 4 equal intervals of 8 frames equal_intervals = torch.chunk(chunks, 4, dim=1) # Step 2: Randomly sample 3 points within the chunk to determine new intervals success = 0 all_success = [] #sample_points = torch.sort(torch.randint(1, 32, (bs, 3), device=chunks.device), dim=1).values # new_intervals_boundaries = torch.cat([torch.zeros((bs, 1), device=chunks.device, dtype=torch.long), sample_points, 32*torch.ones((bs, 1), device=chunks.device, dtype=torch.long)], dim=1) while success != 1: sample_points = sorted(random.sample(range(1, 32), 3)) new_intervals_boundaries = [0] + sample_points + [32] new_intervals = [chunks[0][new_intervals_boundaries[i]:new_intervals_boundaries[i+1]] for i in range(4)] speed_ratios = [8 / len(new_interval) for new_interval in new_intervals] # if any of the speed ratios is greater than 3 or less than 0.33, resample if all([0.33 <= speed_ratio <= 3 for speed_ratio in speed_ratios]): success += 1 all_success.append(new_intervals_boundaries) new_intervals_boundaries = torch.from_numpy(np.array(all_success)) # print(new_intervals_boundaries) all_shapes = new_intervals_boundaries[:, 1:] - new_intervals_boundaries[:, :-1] # Step 4: Adjust the speed of each new interval adjusted_intervals = [] # print(equal_intervals[0].shape) for i in range(4): adjusted_interval = resize_motion_sequence_tensor(equal_intervals[i], all_shapes[0, i]) adjusted_intervals.append(adjusted_interval) # Step 5: Concatenate the adjusted intervals adjusted_chunk = torch.cat(adjusted_intervals, dim=1) return adjusted_chunk def compute_exact_iou(bbox1, bbox2): x1 = max(bbox1[0], bbox2[0]) y1 = max(bbox1[1], bbox2[1]) x2 = min(bbox1[0] + bbox1[2], bbox2[0] + bbox2[2]) y2 = min(bbox1[1] + bbox1[3], bbox2[1] + bbox2[3]) intersection_area = max(0, x2 - x1) * max(0, y2 - y1) bbox1_area = bbox1[2] * bbox1[3] bbox2_area = bbox2[2] * bbox2[3] union_area = bbox1_area + bbox2_area - intersection_area if union_area == 0: return 0 return intersection_area / union_area def compute_iou(mask1, mask2): # Compute the intersection intersection = np.logical_and(mask1, mask2).sum() # Compute the union union = np.logical_or(mask1, mask2).sum() # Compute the IoU iou = intersection / union return iou def blankblending(all_frames, x, n): return all_frames[x:x+n+1] def synthesize_intermediate_frames_FILM(frame1, frame2, t, name, save_path): import replicate from urllib.request import urlretrieve import os cv2.imwrite(save_path[:-9]+name+"_frame1.png", frame1) cv2.imwrite(save_path[:-9]+name+"_frame2.png", frame2) os.environ["REPLICATE_API_TOKEN"] = "r8_He1rkPk9GAxNQ3LpOohK8sYw1SUfMYV3Fxk9b" output = replicate.run( "google-research/frame-interpolation:4f88a16a13673a8b589c18866e540556170a5bcb2ccdc12de556e800e9456d3d", input={ "frame1": open(save_path[:-9]+name+"_frame1.png", "rb"), "frame2": open(save_path[:-9]+name+"_frame2.png", "rb"), "times_to_interpolate": t, } ) print(output) urlretrieve(output, save_path[:-9]+name+"_inter.mp4") return load_video_as_numpy_array(save_path[:-9]+name+"_inter.mp4") def load_video_as_numpy_array(video_path): cap = cv2.VideoCapture(video_path) # Using list comprehension to read frames and store in a list frames = [frame for ret, frame in iter(lambda: cap.read(), (False, None)) if ret] cap.release() return np.array(frames) def synthesize_intermediate_frames_bidirectional(all_frames, x, n): frame1 = all_frames[x] frame2 = all_frames[x + n] # Convert the frames to grayscale gray1 = cv2.cvtColor(frame1, cv2.COLOR_BGR2GRAY) gray2 = cv2.cvtColor(frame2, cv2.COLOR_BGR2GRAY) # Calculate the forward and backward optical flow forward_flow = cv2.calcOpticalFlowFarneback(gray1, gray2, None, 0.5, 3, 15, 3, 5, 1.2, 0) backward_flow = cv2.calcOpticalFlowFarneback(gray2, gray1, None, 0.5, 3, 15, 3, 5, 1.2, 0) synthesized_frames = [] for i in range(1, n): # For each intermediate frame between x and x + n alpha = i / n # Interpolation factor # Compute the intermediate forward and backward flow intermediate_forward_flow = forward_flow * alpha intermediate_backward_flow = backward_flow * (1 - alpha) # Warp the frames based on the intermediate flow h, w = frame1.shape[:2] flow_map = np.column_stack((np.repeat(np.arange(h), w), np.tile(np.arange(w), h))) forward_displacement = flow_map + intermediate_forward_flow.reshape(-1, 2) backward_displacement = flow_map - intermediate_backward_flow.reshape(-1, 2) # Use cv2.remap for efficient warping remap_x_forward, remap_y_forward = np.clip(forward_displacement[:, 1], 0, w - 1), np.clip(forward_displacement[:, 0], 0, h - 1) remap_x_backward, remap_y_backward = np.clip(backward_displacement[:, 1], 0, w - 1), np.clip(backward_displacement[:, 0], 0, h - 1) warped_forward = cv2.remap(frame1, remap_x_forward.reshape(h, w).astype(np.float32), remap_y_forward.reshape(h, w).astype(np.float32), interpolation=cv2.INTER_LINEAR) warped_backward = cv2.remap(frame2, remap_x_backward.reshape(h, w).astype(np.float32), remap_y_backward.reshape(h, w).astype(np.float32), interpolation=cv2.INTER_LINEAR) # Blend the warped frames to generate the intermediate frame intermediate_frame = cv2.addWeighted(warped_forward, 1 - alpha, warped_backward, alpha, 0) synthesized_frames.append(intermediate_frame) return synthesized_frames # Return n-2 synthesized intermediate frames def linear_interpolate_frames(all_frames, x, n): frame1 = all_frames[x] frame2 = all_frames[x + n] synthesized_frames = [] for i in range(1, n): # For each intermediate frame between x and x + n alpha = i / (n) # Correct interpolation factor inter_frame = cv2.addWeighted(frame1, 1 - alpha, frame2, alpha, 0) synthesized_frames.append(inter_frame) return synthesized_frames[:-1] def warp_frame(src_frame, flow): h, w = flow.shape[:2] flow_map = np.column_stack((np.repeat(np.arange(h), w), np.tile(np.arange(w), h))) displacement = flow_map + flow.reshape(-1, 2) # Extract x and y coordinates of the displacement x_coords = np.clip(displacement[:, 1], 0, w - 1).reshape(h, w).astype(np.float32) y_coords = np.clip(displacement[:, 0], 0, h - 1).reshape(h, w).astype(np.float32) # Use cv2.remap for efficient warping warped_frame = cv2.remap(src_frame, x_coords, y_coords, interpolation=cv2.INTER_LINEAR) return warped_frame def synthesize_intermediate_frames(all_frames, x, n): # Calculate Optical Flow between the first and last frame frame1 = cv2.cvtColor(all_frames[x], cv2.COLOR_BGR2GRAY) frame2 = cv2.cvtColor(all_frames[x + n], cv2.COLOR_BGR2GRAY) flow = cv2.calcOpticalFlowFarneback(frame1, frame2, None, 0.5, 3, 15, 3, 5, 1.2, 0) synthesized_frames = [] for i in range(1, n): # For each intermediate frame alpha = i / (n) # Interpolation factor intermediate_flow = flow * alpha # Interpolate the flow intermediate_frame = warp_frame(all_frames[x], intermediate_flow) # Warp the first frame synthesized_frames.append(intermediate_frame) return synthesized_frames def map2color(s): m = hashlib.md5() m.update(s.encode('utf-8')) color_code = m.hexdigest()[:6] return '#' + color_code def euclidean_distance(a, b): return np.sqrt(np.sum((a - b)**2)) def adjust_array(x, k): len_x = len(x) len_k = len(k) # If x is shorter than k, pad with zeros if len_x < len_k: return np.pad(x, (0, len_k - len_x), 'constant') # If x is longer than k, truncate x elif len_x > len_k: return x[:len_k] # If both are of same length else: return x def onset_to_frame(onset_times, audio_length, fps): # Calculate total number of frames for the given audio length total_frames = int(audio_length * fps) # Create an array of zeros of shape (total_frames,) frame_array = np.zeros(total_frames, dtype=np.int32) # For each onset time, calculate the frame number and set it to 1 for onset in onset_times: frame_num = int(onset * fps) # Check if the frame number is within the array bounds if 0 <= frame_num < total_frames: frame_array[frame_num] = 1 return frame_array # def np_slerp(q1, q2, t): # dot_product = np.sum(q1 * q2, axis=-1) # q2_flip = np.where(dot_product[:, None] < 0, -q2, q2) # Flip quaternions where dot_product is negative # dot_product = np.abs(dot_product) # angle = np.arccos(np.clip(dot_product, -1, 1)) # sin_angle = np.sin(angle) # t1 = np.sin((1.0 - t) * angle) / sin_angle # t2 = np.sin(t * angle) / sin_angle # return t1 * q1 + t2 * q2_flip def smooth_rotvec_animations(animation1, animation2, blend_frames): """ Smoothly transition between two animation clips using SLERP. Parameters: - animation1: The first animation clip, a numpy array of shape [n, k]. - animation2: The second animation clip, a numpy array of shape [n, k]. - blend_frames: Number of frames over which to blend the two animations. Returns: - A smoothly blended animation clip of shape [2n, k]. """ # Ensure blend_frames doesn't exceed the length of either animation n1, k1 = animation1.shape n2, k2 = animation2.shape animation1 = animation1.reshape(n1, k1//3, 3) animation2 = animation2.reshape(n2, k2//3, 3) blend_frames = min(blend_frames, len(animation1), len(animation2)) all_int = [] for i in range(k1//3): # Convert rotation vectors to quaternion for the overlapping part q = R.from_rotvec(np.concatenate([animation1[0:1, i], animation2[-2:-1, i]], axis=0))#.as_quat() # q2 = R.from_rotvec()#.as_quat() times = [0, blend_frames * 2 - 1] slerp = Slerp(times, q) interpolated = slerp(np.arange(blend_frames * 2)) interpolated_rotvecs = interpolated.as_rotvec() all_int.append(interpolated_rotvecs) interpolated_rotvecs = np.concatenate(all_int, axis=1) # result = np.vstack((animation1[:-blend_frames], interpolated_rotvecs, animation2[blend_frames:])) result = interpolated_rotvecs.reshape(2*n1, k1) return result def smooth_animations(animation1, animation2, blend_frames): """ Smoothly transition between two animation clips using linear interpolation. Parameters: - animation1: The first animation clip, a numpy array of shape [n, k]. - animation2: The second animation clip, a numpy array of shape [n, k]. - blend_frames: Number of frames over which to blend the two animations. Returns: - A smoothly blended animation clip of shape [2n, k]. """ # Ensure blend_frames doesn't exceed the length of either animation blend_frames = min(blend_frames, len(animation1), len(animation2)) # Extract overlapping sections overlap_a1 = animation1[-blend_frames:-blend_frames+1, :] overlap_a2 = animation2[blend_frames-1:blend_frames, :] # Create blend weights for linear interpolation alpha = np.linspace(0, 1, 2 * blend_frames).reshape(-1, 1) # Linearly interpolate between overlapping sections blended_overlap = overlap_a1 * (1 - alpha) + overlap_a2 * alpha # Extend the animations to form the result with 2n frames if blend_frames == len(animation1) and blend_frames == len(animation2): result = blended_overlap else: before_blend = animation1[:-blend_frames] after_blend = animation2[blend_frames:] result = np.vstack((before_blend, blended_overlap, after_blend)) return result def interpolate_sequence(quaternions): bs, n, j, _ = quaternions.shape new_n = 2 * n new_quaternions = torch.zeros((bs, new_n, j, 4), device=quaternions.device, dtype=quaternions.dtype) for i in range(n): q1 = quaternions[:, i, :, :] new_quaternions[:, 2*i, :, :] = q1 if i < n - 1: q2 = quaternions[:, i + 1, :, :] new_quaternions[:, 2*i + 1, :, :] = slerp(q1, q2, 0.5) else: # For the last point, duplicate the value new_quaternions[:, 2*i + 1, :, :] = q1 return new_quaternions def quaternion_multiply(q1, q2): w1, x1, y1, z1 = q1 w2, x2, y2, z2 = q2 w = w1 * w2 - x1 * x2 - y1 * y2 - z1 * z2 x = w1 * x2 + x1 * w2 + y1 * z2 - z1 * y2 y = w1 * y2 + y1 * w2 + z1 * x2 - x1 * z2 z = w1 * z2 + z1 * w2 + x1 * y2 - y1 * x2 return w, x, y, z def quaternion_conjugate(q): w, x, y, z = q return (w, -x, -y, -z) def slerp(q1, q2, t): dot = torch.sum(q1 * q2, dim=-1, keepdim=True) flip = (dot < 0).float() q2 = (1 - flip * 2) * q2 dot = dot * (1 - flip * 2) DOT_THRESHOLD = 0.9995 mask = (dot > DOT_THRESHOLD).float() theta_0 = torch.acos(dot) theta = theta_0 * t q3 = q2 - q1 * dot q3 = q3 / torch.norm(q3, dim=-1, keepdim=True) interpolated = (torch.cos(theta) * q1 + torch.sin(theta) * q3) return mask * (q1 + t * (q2 - q1)) + (1 - mask) * interpolated def estimate_linear_velocity(data_seq, dt): ''' Given some batched data sequences of T timesteps in the shape (B, T, ...), estimates the velocity for the middle T-2 steps using a second order central difference scheme. The first and last frames are with forward and backward first-order differences, respectively - h : step size ''' # first steps is forward diff (t+1 - t) / dt init_vel = (data_seq[:, 1:2] - data_seq[:, :1]) / dt # middle steps are second order (t+1 - t-1) / 2dt middle_vel = (data_seq[:, 2:] - data_seq[:, 0:-2]) / (2 * dt) # last step is backward diff (t - t-1) / dt final_vel = (data_seq[:, -1:] - data_seq[:, -2:-1]) / dt vel_seq = torch.cat([init_vel, middle_vel, final_vel], dim=1) return vel_seq def velocity2position(data_seq, dt, init_pos): res_trans = [] for i in range(data_seq.shape[1]): if i == 0: res_trans.append(init_pos.unsqueeze(1)) else: res = data_seq[:, i-1:i] * dt + res_trans[-1] res_trans.append(res) return torch.cat(res_trans, dim=1) def estimate_angular_velocity(rot_seq, dt): ''' Given a batch of sequences of T rotation matrices, estimates angular velocity at T-2 steps. Input sequence should be of shape (B, T, ..., 3, 3) ''' # see https://en.wikipedia.org/wiki/Angular_velocity#Calculation_from_the_orientation_matrix dRdt = estimate_linear_velocity(rot_seq, dt) R = rot_seq RT = R.transpose(-1, -2) # compute skew-symmetric angular velocity tensor w_mat = torch.matmul(dRdt, RT) # pull out angular velocity vector by averaging symmetric entries w_x = (-w_mat[..., 1, 2] + w_mat[..., 2, 1]) / 2.0 w_y = (w_mat[..., 0, 2] - w_mat[..., 2, 0]) / 2.0 w_z = (-w_mat[..., 0, 1] + w_mat[..., 1, 0]) / 2.0 w = torch.stack([w_x, w_y, w_z], axis=-1) return w def image_from_bytes(image_bytes): import matplotlib.image as mpimg from io import BytesIO return mpimg.imread(BytesIO(image_bytes), format='PNG') def process_frame(i, vertices_all, vertices1_all, faces, output_dir, filenames): import matplotlib matplotlib.use('Agg') import matplotlib.pyplot as plt import trimesh import pyrender def deg_to_rad(degrees): return degrees * np.pi / 180 uniform_color = [220, 220, 220, 255] resolution = (1000, 1000) figsize = (10, 10) fig, axs = plt.subplots( nrows=1, ncols=2, figsize=(figsize[0] * 2, figsize[1] * 1) ) axs = axs.flatten() vertices = vertices_all[i] vertices1 = vertices1_all[i] filename = f"{output_dir}frame_{i}.png" filenames.append(filename) if i%100 == 0: print('processed', i, 'frames') #time_s = time.time() #print(vertices.shape) angle_rad = deg_to_rad(-2) pose_camera = np.array([ [1.0, 0.0, 0.0, 0.0], [0.0, np.cos(angle_rad), -np.sin(angle_rad), 1.0], [0.0, np.sin(angle_rad), np.cos(angle_rad), 5.0], [0.0, 0.0, 0.0, 1.0] ]) angle_rad = deg_to_rad(-30) pose_light = np.array([ [1.0, 0.0, 0.0, 0.0], [0.0, np.cos(angle_rad), -np.sin(angle_rad), 0.0], [0.0, np.sin(angle_rad), np.cos(angle_rad), 3.0], [0.0, 0.0, 0.0, 1.0] ]) for vtx_idx, vtx in enumerate([vertices, vertices1]): trimesh_mesh = trimesh.Trimesh( vertices=vtx, faces=faces, vertex_colors=uniform_color ) mesh = pyrender.Mesh.from_trimesh( trimesh_mesh, smooth=True ) scene = pyrender.Scene() scene.add(mesh) camera = pyrender.OrthographicCamera(xmag=1.0, ymag=1.0) scene.add(camera, pose=pose_camera) light = pyrender.DirectionalLight(color=[1.0, 1.0, 1.0], intensity=4.0) scene.add(light, pose=pose_light) renderer = pyrender.OffscreenRenderer(*resolution) color, _ = renderer.render(scene) axs[vtx_idx].imshow(color) axs[vtx_idx].axis('off') renderer.delete() plt.savefig(filename, bbox_inches='tight') plt.close(fig) def generate_images(frames, vertices_all, vertices1_all, faces, output_dir, filenames): import multiprocessing # import trimesh num_cores = multiprocessing.cpu_count() - 1 # This will get the number of cores on your machine. # mesh = trimesh.Trimesh(vertices_all[0], faces) # scene = mesh.scene() # fov = scene.camera.fov.copy() # fov[0] = 80.0 # fov[1] = 60.0 # camera_params = { # 'fov': fov, # 'resolution': scene.camera.resolution, # 'focal': scene.camera.focal, # 'z_near': scene.camera.z_near, # "z_far": scene.camera.z_far, # 'transform': scene.graph[scene.camera.name][0] # } # mesh1 = trimesh.Trimesh(vertices1_all[0], faces) # scene1 = mesh1.scene() # camera_params1 = { # 'fov': fov, # 'resolution': scene1.camera.resolution, # 'focal': scene1.camera.focal, # 'z_near': scene1.camera.z_near, # "z_far": scene1.camera.z_far, # 'transform': scene1.graph[scene1.camera.name][0] # } # Use a Pool to manage the processes # print(num_cores) # for i in range(frames): # process_frame(i, vertices_all, vertices1_all, faces, output_dir, use_matplotlib, filenames, camera_params, camera_params1) for i in range(frames): process_frame(i*3, vertices_all, vertices1_all, faces, output_dir, filenames) # progress = multiprocessing.Value('i', 0) # lock = multiprocessing.Lock() # with multiprocessing.Pool(num_cores) as pool: # # pool.starmap(process_frame, [(i, vertices_all, vertices1_all, faces, output_dir, use_matplotlib, filenames, camera_params, camera_params1) for i in range(frames)]) # pool.starmap( # process_frame, # [ # (i, vertices_all, vertices1_all, faces, output_dir, filenames) # for i in range(frames) # ] # ) # progress = multiprocessing.Value('i', 0) # lock = multiprocessing.Lock() # with multiprocessing.Pool(num_cores) as pool: # # pool.starmap(process_frame, [(i, vertices_all, vertices1_all, faces, output_dir, use_matplotlib, filenames, camera_params, camera_params1) for i in range(frames)]) # pool.starmap( # process_frame, # [ # (i, vertices_all, vertices1_all, faces, output_dir, filenames) # for i in range(frames) # ] # ) def render_one_sequence_with_face( res_npz_path, gt_npz_path, output_dir, audio_path, model_folder="/data/datasets/smplx_models/", model_type='smplx', gender='NEUTRAL_2020', ext='npz', num_betas=300, num_expression_coeffs=100, use_face_contour=False, use_matplotlib=False, args=None): import smplx import matplotlib.pyplot as plt import imageio from tqdm import tqdm import os import numpy as np import torch import moviepy.editor as mp import librosa model = smplx.create(model_folder, model_type=model_type, gender=gender, use_face_contour=use_face_contour, num_betas=num_betas, num_expression_coeffs=num_expression_coeffs, ext=ext, use_pca=False).cuda() #data_npz = np.load(f"{output_dir}{res_npz_path}.npz") data_np_body = np.load(res_npz_path, allow_pickle=True) gt_np_body = np.load(gt_npz_path, allow_pickle=True) if not os.path.exists(output_dir): os.makedirs(output_dir) # if not use_matplotlib: # import trimesh #import pyrender from pyvirtualdisplay import Display #''' #display = Display(visible=0, size=(1000, 1000)) #display.start() faces = np.load(f"{model_folder}/smplx/SMPLX_NEUTRAL_2020.npz", allow_pickle=True)["f"] seconds = 1 #data_npz["jaw_pose"].shape[0] n = data_np_body["poses"].shape[0] beta = torch.from_numpy(data_np_body["betas"]).to(torch.float32).unsqueeze(0).cuda() beta = beta.repeat(n, 1) expression = torch.from_numpy(data_np_body["expressions"][:n]).to(torch.float32).cuda() jaw_pose = torch.from_numpy(data_np_body["poses"][:n, 66:69]).to(torch.float32).cuda() pose = torch.from_numpy(data_np_body["poses"][:n]).to(torch.float32).cuda() transl = torch.from_numpy(data_np_body["trans"][:n]).to(torch.float32).cuda() # print(beta.shape, expression.shape, jaw_pose.shape, pose.shape, transl.shape, pose[:,:3].shape) output = model(betas=beta, transl=transl, expression=expression, jaw_pose=jaw_pose, global_orient=pose[:,:3], body_pose=pose[:,3:21*3+3], left_hand_pose=pose[:,25*3:40*3], right_hand_pose=pose[:,40*3:55*3], leye_pose=pose[:, 69:72], reye_pose=pose[:, 72:75], return_verts=True) vertices_all = output["vertices"].cpu().detach().numpy() beta1 = torch.from_numpy(data_np_body["betas"]).to(torch.float32).unsqueeze(0).cuda() beta1 = beta1.repeat(n, 1) expression1 = torch.from_numpy(data_np_body["expressions"][:n]).to(torch.float32).cuda() zero_pose = np.zeros_like(data_np_body["poses"]) jaw_pose1 = torch.from_numpy(zero_pose[:n,66:69]).to(torch.float32).cuda() pose1 = torch.from_numpy(zero_pose[:n]).to(torch.float32).cuda() zero_trans = np.zeros_like(data_np_body["trans"]) transl1 = torch.from_numpy(zero_trans[:n]).to(torch.float32).cuda() output1 = model(betas=beta1, transl=transl1, expression=expression1, jaw_pose=jaw_pose1, global_orient=pose1[:,:3], body_pose=pose1[:,3:21*3+3], left_hand_pose=pose1[:,25*3:40*3], right_hand_pose=pose1[:,40*3:55*3], leye_pose=pose1[:, 69:72], reye_pose=pose1[:, 72:75], return_verts=True) vertices1_all = output1["vertices"].cpu().detach().numpy()*8 trans_down = np.zeros_like(vertices1_all) trans_down[:, :, 1] = 1.55 vertices1_all = vertices1_all - trans_down if args.debug: seconds = 1 else: seconds = vertices_all.shape[0]//30 silent_video_file_path = utils.fast_render.generate_silent_videos(args.render_video_fps, args.render_video_width, args.render_video_height, args.render_concurrent_num, args.render_tmp_img_filetype, int(seconds*args.render_video_fps), vertices1_all, vertices_all, faces, output_dir) base_filename_without_ext = os.path.splitext(os.path.basename(res_npz_path))[0] final_clip = os.path.join(output_dir, f"{base_filename_without_ext}.mp4") utils.media.add_audio_to_video(silent_video_file_path, audio_path, final_clip) os.remove(silent_video_file_path) return final_clip def render_one_sequence( res_npz_path, gt_npz_path, output_dir, audio_path, model_folder="/data/datasets/smplx_models/", model_type='smplx', gender='NEUTRAL_2020', ext='npz', num_betas=300, num_expression_coeffs=100, use_face_contour=False, use_matplotlib=False, args=None): import smplx import matplotlib.pyplot as plt import imageio from tqdm import tqdm import os import numpy as np import torch import moviepy.editor as mp import librosa model = smplx.create(model_folder, model_type=model_type, gender=gender, use_face_contour=use_face_contour, num_betas=num_betas, num_expression_coeffs=num_expression_coeffs, ext=ext, use_pca=False).cuda() #data_npz = np.load(f"{output_dir}{res_npz_path}.npz") data_np_body = np.load(res_npz_path, allow_pickle=True) gt_np_body = np.load(gt_npz_path, allow_pickle=True) if not os.path.exists(output_dir): os.makedirs(output_dir) # if not use_matplotlib: # import trimesh #import pyrender from pyvirtualdisplay import Display #''' #display = Display(visible=0, size=(1000, 1000)) #display.start() faces = np.load(f"{model_folder}/smplx/SMPLX_NEUTRAL_2020.npz", allow_pickle=True)["f"] seconds = 1 #data_npz["jaw_pose"].shape[0] n = data_np_body["poses"].shape[0] beta = torch.from_numpy(data_np_body["betas"]).to(torch.float32).unsqueeze(0).cuda() beta = beta.repeat(n, 1) expression = torch.from_numpy(data_np_body["expressions"][:n]).to(torch.float32).cuda() jaw_pose = torch.from_numpy(data_np_body["poses"][:n, 66:69]).to(torch.float32).cuda() pose = torch.from_numpy(data_np_body["poses"][:n]).to(torch.float32).cuda() transl = torch.from_numpy(data_np_body["trans"][:n]).to(torch.float32).cuda() # print(beta.shape, expression.shape, jaw_pose.shape, pose.shape, transl.shape, pose[:,:3].shape) output = model(betas=beta, transl=transl, expression=expression, jaw_pose=jaw_pose, global_orient=pose[:,:3], body_pose=pose[:,3:21*3+3], left_hand_pose=pose[:,25*3:40*3], right_hand_pose=pose[:,40*3:55*3], leye_pose=pose[:, 69:72], reye_pose=pose[:, 72:75], return_verts=True) vertices_all = output["vertices"].cpu().detach().numpy() beta1 = torch.from_numpy(gt_np_body["betas"]).to(torch.float32).unsqueeze(0).cuda() expression1 = torch.from_numpy(gt_np_body["expressions"][:n]).to(torch.float32).cuda() jaw_pose1 = torch.from_numpy(gt_np_body["poses"][:n,66:69]).to(torch.float32).cuda() pose1 = torch.from_numpy(gt_np_body["poses"][:n]).to(torch.float32).cuda() transl1 = torch.from_numpy(gt_np_body["trans"][:n]).to(torch.float32).cuda() output1 = model(betas=beta1, transl=transl1, expression=expression1, jaw_pose=jaw_pose1, global_orient=pose1[:,:3], body_pose=pose1[:,3:21*3+3], left_hand_pose=pose1[:,25*3:40*3], right_hand_pose=pose1[:,40*3:55*3], leye_pose=pose1[:, 69:72], reye_pose=pose1[:, 72:75],return_verts=True) vertices1_all = output1["vertices"].cpu().detach().numpy() if args.debug: seconds = 1 else: seconds = vertices_all.shape[0]//30 silent_video_file_path = utils.fast_render.generate_silent_videos(args.render_video_fps, args.render_video_width, args.render_video_height, args.render_concurrent_num, args.render_tmp_img_filetype, int(seconds*args.render_video_fps), vertices_all, vertices1_all, faces, output_dir) base_filename_without_ext = os.path.splitext(os.path.basename(res_npz_path))[0] final_clip = os.path.join(output_dir, f"{base_filename_without_ext}.mp4") utils.media.add_audio_to_video(silent_video_file_path, audio_path, final_clip) os.remove(silent_video_file_path) return final_clip def render_one_sequence_no_gt( res_npz_path, output_dir, audio_path, model_folder="/data/datasets/smplx_models/", model_type='smplx', gender='NEUTRAL_2020', ext='npz', num_betas=300, num_expression_coeffs=100, use_face_contour=False, use_matplotlib=False, args=None): import smplx import matplotlib.pyplot as plt import imageio from tqdm import tqdm import os import numpy as np import torch import moviepy.editor as mp import librosa model = smplx.create(model_folder, model_type=model_type, gender=gender, use_face_contour=use_face_contour, num_betas=num_betas, num_expression_coeffs=num_expression_coeffs, ext=ext, use_pca=False).cuda() #data_npz = np.load(f"{output_dir}{res_npz_path}.npz") data_np_body = np.load(res_npz_path, allow_pickle=True) # gt_np_body = np.load(gt_npz_path, allow_pickle=True) if not os.path.exists(output_dir): os.makedirs(output_dir) # if not use_matplotlib: # import trimesh #import pyrender from pyvirtualdisplay import Display #''' #display = Display(visible=0, size=(1000, 1000)) #display.start() faces = np.load(f"{model_folder}/smplx/SMPLX_NEUTRAL_2020.npz", allow_pickle=True)["f"] seconds = 1 #data_npz["jaw_pose"].shape[0] n = data_np_body["poses"].shape[0] beta = torch.from_numpy(data_np_body["betas"]).to(torch.float32).unsqueeze(0).cuda() beta = beta.repeat(n, 1) expression = torch.from_numpy(data_np_body["expressions"][:n]).to(torch.float32).cuda() jaw_pose = torch.from_numpy(data_np_body["poses"][:n, 66:69]).to(torch.float32).cuda() pose = torch.from_numpy(data_np_body["poses"][:n]).to(torch.float32).cuda() transl = torch.from_numpy(data_np_body["trans"][:n]).to(torch.float32).cuda() # print(beta.shape, expression.shape, jaw_pose.shape, pose.shape, transl.shape, pose[:,:3].shape) output = model(betas=beta, transl=transl, expression=expression, jaw_pose=jaw_pose, global_orient=pose[:,:3], body_pose=pose[:,3:21*3+3], left_hand_pose=pose[:,25*3:40*3], right_hand_pose=pose[:,40*3:55*3], leye_pose=pose[:, 69:72], reye_pose=pose[:, 72:75], return_verts=True) vertices_all = output["vertices"].cpu().detach().numpy() # beta1 = torch.from_numpy(gt_np_body["betas"]).to(torch.float32).unsqueeze(0).cuda() # expression1 = torch.from_numpy(gt_np_body["expressions"][:n]).to(torch.float32).cuda() # jaw_pose1 = torch.from_numpy(gt_np_body["poses"][:n,66:69]).to(torch.float32).cuda() # pose1 = torch.from_numpy(gt_np_body["poses"][:n]).to(torch.float32).cuda() # transl1 = torch.from_numpy(gt_np_body["trans"][:n]).to(torch.float32).cuda() # output1 = model(betas=beta1, transl=transl1, expression=expression1, jaw_pose=jaw_pose1, global_orient=pose1[:,:3], body_pose=pose1[:,3:21*3+3], left_hand_pose=pose1[:,25*3:40*3], right_hand_pose=pose1[:,40*3:55*3], # leye_pose=pose1[:, 69:72], # reye_pose=pose1[:, 72:75],return_verts=True) # vertices1_all = output1["vertices"].cpu().detach().numpy() if args.debug: seconds = 1 else: seconds = vertices_all.shape[0]//30 silent_video_file_path = utils.fast_render.generate_silent_videos_no_gt(args.render_video_fps, args.render_video_width, args.render_video_height, args.render_concurrent_num, args.render_tmp_img_filetype, int(seconds*args.render_video_fps), vertices_all, faces, output_dir) base_filename_without_ext = os.path.splitext(os.path.basename(res_npz_path))[0] final_clip = os.path.join(output_dir, f"{base_filename_without_ext}.mp4") utils.media.add_audio_to_video(silent_video_file_path, audio_path, final_clip) os.remove(silent_video_file_path) return final_clip def print_exp_info(args): logger.info(pprint.pformat(vars(args))) logger.info(f"# ------------ {args.name} ----------- #") logger.info("PyTorch version: {}".format(torch.__version__)) logger.info("CUDA version: {}".format(torch.version.cuda)) logger.info("{} GPUs".format(torch.cuda.device_count())) logger.info(f"Random Seed: {args.random_seed}") def args2csv(args, get_head=False, list4print=[]): for k, v in args.items(): if isinstance(args[k], dict): args2csv(args[k], get_head, list4print) else: list4print.append(k) if get_head else list4print.append(v) return list4print class EpochTracker: def __init__(self, metric_names, metric_directions): assert len(metric_names) == len(metric_directions), "Metric names and directions should have the same length" self.metric_names = metric_names self.states = ['train', 'val', 'test'] self.types = ['last', 'best'] self.values = {name: {state: {type_: {'value': np.inf if not is_higher_better else -np.inf, 'epoch': 0} for type_ in self.types} for state in self.states} for name, is_higher_better in zip(metric_names, metric_directions)} self.loss_meters = {name: {state: AverageMeter(f"{name}_{state}") for state in self.states} for name in metric_names} self.is_higher_better = {name: direction for name, direction in zip(metric_names, metric_directions)} self.train_history = {name: [] for name in metric_names} self.val_history = {name: [] for name in metric_names} def update_meter(self, name, state, value): self.loss_meters[name][state].update(value) def update_values(self, name, state, epoch): value_avg = self.loss_meters[name][state].avg new_best = False if ((value_avg < self.values[name][state]['best']['value'] and not self.is_higher_better[name]) or (value_avg > self.values[name][state]['best']['value'] and self.is_higher_better[name])): self.values[name][state]['best']['value'] = value_avg self.values[name][state]['best']['epoch'] = epoch new_best = True self.values[name][state]['last']['value'] = value_avg self.values[name][state]['last']['epoch'] = epoch return new_best def get(self, name, state, type_): return self.values[name][state][type_] def reset(self): for name in self.metric_names: for state in self.states: self.loss_meters[name][state].reset() def flatten_values(self): flat_dict = {} for name in self.metric_names: for state in self.states: for type_ in self.types: value_key = f"{name}_{state}_{type_}" epoch_key = f"{name}_{state}_{type_}_epoch" flat_dict[value_key] = self.values[name][state][type_]['value'] flat_dict[epoch_key] = self.values[name][state][type_]['epoch'] return flat_dict def update_and_plot(self, name, epoch, save_path): new_best_train = self.update_values(name, 'train', epoch) new_best_val = self.update_values(name, 'val', epoch) self.train_history[name].append(self.loss_meters[name]['train'].avg) self.val_history[name].append(self.loss_meters[name]['val'].avg) train_values = self.train_history[name] val_values = self.val_history[name] epochs = list(range(1, len(train_values) + 1)) plt.figure(figsize=(10, 6)) plt.plot(epochs, train_values, label='Train') plt.plot(epochs, val_values, label='Val') plt.title(f'Train vs Val {name} over epochs') plt.xlabel('Epochs') plt.ylabel(name) plt.legend() plt.savefig(save_path) plt.close() return new_best_train, new_best_val def record_trial(args, tracker): """ 1. record notes, score, env_name, experments_path, """ csv_path = args.out_path + "custom/" +args.csv_name+".csv" all_print_dict = vars(args) all_print_dict.update(tracker.flatten_values()) if not os.path.exists(csv_path): pd.DataFrame([all_print_dict]).to_csv(csv_path, index=False) else: df_existing = pd.read_csv(csv_path) df_new = pd.DataFrame([all_print_dict]) df_aligned = df_existing.append(df_new).fillna("") df_aligned.to_csv(csv_path, index=False) def set_random_seed(args): os.environ['PYTHONHASHSEED'] = str(args.random_seed) random.seed(args.random_seed) np.random.seed(args.random_seed) torch.manual_seed(args.random_seed) torch.cuda.manual_seed_all(args.random_seed) torch.cuda.manual_seed(args.random_seed) torch.backends.cudnn.deterministic = args.deterministic #args.CUDNN_DETERMINISTIC torch.backends.cudnn.benchmark = args.benchmark torch.backends.cudnn.enabled = args.cudnn_enabled def save_checkpoints(save_path, model, opt=None, epoch=None, lrs=None): if lrs is not None: states = { 'model_state': model.state_dict(), 'epoch': epoch + 1, 'opt_state': opt.state_dict(), 'lrs':lrs.state_dict(),} elif opt is not None: states = { 'model_state': model.state_dict(), 'epoch': epoch + 1, 'opt_state': opt.state_dict(),} else: states = { 'model_state': model.state_dict(),} torch.save(states, save_path) def load_checkpoints(model, save_path, load_name='model'): states = torch.load(save_path) new_weights = OrderedDict() flag=False for k, v in states['model_state'].items(): #print(k) if "module" not in k: break else: new_weights[k[7:]]=v flag=True if flag: try: model.load_state_dict(new_weights) except: #print(states['model_state']) model.load_state_dict(states['model_state']) else: model.load_state_dict(states['model_state']) logger.info(f"load self-pretrained checkpoints for {load_name}") def model_complexity(model, args): from ptflops import get_model_complexity_info flops, params = get_model_complexity_info(model, (args.T_GLOBAL._DIM, args.TRAIN.CROP, args.TRAIN), as_strings=False, print_per_layer_stat=False) logging.info('{:<30} {:<8} BFlops'.format('Computational complexity: ', flops / 1e9)) logging.info('{:<30} {:<8} MParams'.format('Number of parameters: ', params / 1e6)) class AverageMeter(object): """Computes and stores the average and current value""" def __init__(self, name, fmt=':f'): self.name = name self.fmt = fmt self.reset() def reset(self): self.val = 0 self.avg = 0 self.sum = 0 self.count = 0 def update(self, val, n=1): self.val = val self.sum += val * n self.count += n self.avg = self.sum / self.count def __str__(self): fmtstr = '{name} {val' + self.fmt + '} ({avg' + self.fmt + '})' return fmtstr.format(**self.__dict__)