latte2512/TANGO
0
1# import smplx2# import torch3# import pickle4# import numpy as np5 6# # Global: Load the SMPL-X model once7# smplx_model = smplx.create(8# "/content/drive/MyDrive/003_Codes/TANGO-JointEmbedding/emage/smplx_models/", 9# model_type='smplx',10# gender='NEUTRAL_2020', 11# use_face_contour=False,12# num_betas=300,13# num_expression_coeffs=100, 14# ext='npz',15# use_pca=True,16# num_pca_comps=12,17# ).to("cuda").eval()18 19# device = "cuda"20 21# def pkl_to_npz(pkl_path, npz_path):22# # Load the pickle file23# with open(pkl_path, "rb") as f:24# pkl_example = pickle.load(f)25 26# bs = 127# n = pkl_example["expression"].shape[0] # Assuming this is the batch size28 29# # Convert numpy arrays to torch tensors30# def to_tensor(numpy_array):31# return torch.tensor(numpy_array, dtype=torch.float32).to(device)32 33# # Ensure that betas are loaded from the pickle data, converting them to torch tensors34# betas = to_tensor(pkl_example["betas"])35# transl = to_tensor(pkl_example["transl"])36# expression = to_tensor(pkl_example["expression"])37# jaw_pose = to_tensor(pkl_example["jaw_pose"])38# global_orient = to_tensor(pkl_example["global_orient"])39# body_pose_axis = to_tensor(pkl_example["body_pose_axis"])40# left_hand_pose = to_tensor(pkl_example['left_hand_pose'])41# right_hand_pose = to_tensor(pkl_example['right_hand_pose'])42# leye_pose = to_tensor(pkl_example['leye_pose'])43# reye_pose = to_tensor(pkl_example['reye_pose'])44 45# # Pass the loaded data into the SMPL-X model46# gt_vertex = smplx_model(47# betas=betas,48# transl=transl, # Translation49# expression=expression, # Expression50# jaw_pose=jaw_pose, # Jaw pose51# global_orient=global_orient, # Global orientation52# body_pose=body_pose_axis, # Body pose53# left_hand_pose=left_hand_pose, # Left hand pose54# right_hand_pose=right_hand_pose, # Right hand pose55# return_full_pose=True,56# leye_pose=leye_pose, # Left eye pose57# reye_pose=reye_pose, # Right eye pose58# )59 60# # Save the relevant data to an npz file61# np.savez(npz_path,62# betas=pkl_example["betas"],63# poses=gt_vertex["full_pose"].cpu().numpy(),64# expressions=pkl_example["expression"],65# trans=pkl_example["transl"],66# model='smplx2020',67# gender='neutral',68# mocap_frame_rate=30,69# )70 71# from tqdm import tqdm72# import os73# def convert_all_pkl_in_folder(folder_path):74# # Collect all .pkl files75# pkl_files = []76# for root, dirs, files in os.walk(folder_path):77# for file in files:78# if file.endswith(".pkl"):79# pkl_files.append(os.path.join(root, file))80 81# # Process each file with a progress bar82# for pkl_path in tqdm(pkl_files, desc="Converting .pkl to .npz"):83# npz_path = pkl_path.replace(".pkl", ".npz") # Replace .pkl with .npz84# pkl_to_npz(pkl_path, npz_path)85 86# convert_all_pkl_in_folder("/content/oliver/oliver/") 87 88 89# import os90# import json91 92# def collect_dataset_info(root_dir):93# dataset_info = []94 95# for root, dirs, files in os.walk(root_dir):96# for file in files:97# if file.endswith(".npz"):98# video_id = file[:-4] # Removing the .npz extension to get the video ID99 100# # Construct the paths based on the current root directory101# motion_path = os.path.join(root)102# video_path = os.path.join(root)103# audio_path = os.path.join(root)104 105# # Determine the mode (train, val, test) by checking parent directory106# mode = root.split(os.sep)[-2] # Assuming mode is one folder up in hierarchy107 108# dataset_info.append({109# "video_id": video_id,110# "video_path": video_path,111# "audio_path": audio_path,112# "motion_path": motion_path,113# "mode": mode 114# })115# return dataset_info116 117# # Set the root directory path of your dataset118# root_dir = '/content/oliver/oliver/' # Adjust this to your actual root directory119# dataset_info = collect_dataset_info(root_dir)120# output_file = '/content/drive/MyDrive/003_Codes/TANGO-JointEmbedding/datasets/show-oliver-original.json'121 122# # Save the dataset information to a JSON file123# with open(output_file, 'w') as json_file:124# json.dump(dataset_info, json_file, indent=4)125# print(f"Dataset information saved to {output_file}")126 127 128# import os129# import json130# import numpy as np131 132# def load_npz(npz_path):133# try:134# data = np.load(npz_path)135# return data136# except Exception as e:137# print(f"Error loading {npz_path}: {e}")138# return None139 140# def generate_clips(data, stride, window_length):141# clips = []142# for entry in data:143# npz_data = load_npz(os.path.join(entry['motion_path'],entry['video_id']+".npz"))144 145# # Only continue if the npz file is successfully loaded146# if npz_data is None:147# continue148 149# # Determine the total length of the sequence from npz data150# total_frames = npz_data["poses"].shape[0]151 152# # Generate clips based on stride and window_length153# for start_idx in range(0, total_frames - window_length + 1, stride):154# end_idx = start_idx + window_length155# clip = {156# "video_id": entry["video_id"],157# "video_path": entry["video_path"],158# "audio_path": entry["audio_path"],159# "motion_path": entry["motion_path"],160# "mode": entry["mode"],161# "start_idx": start_idx,162# "end_idx": end_idx163# }164# clips.append(clip)165 166# return clips167 168# # Load the existing dataset JSON file169# input_json = '/content/drive/MyDrive/003_Codes/TANGO-JointEmbedding/datasets/show-oliver-original.json'170# with open(input_json, 'r') as f:171# dataset_info = json.load(f)172 173# # Set stride and window length174# stride = 40 # Adjust stride as needed175# window_length = 64 # Adjust window length as needed176 177# # Generate clips for all data178# clips_data = generate_clips(dataset_info, stride, window_length)179 180# # Save the filtered clips data to a new JSON file181# output_json = f'/content/drive/MyDrive/003_Codes/TANGO-JointEmbedding/datasets/show-oliver-s{stride}_w{window_length}.json'182# with open(output_json, 'w') as f:183# json.dump(clips_data, f, indent=4)184 185# print(f"Filtered clips data saved to {output_json}")186 187 188 189from ast import Expression190import os191import numpy as np192import wave193from moviepy.editor import VideoFileClip194 195def split_npz(npz_path, output_prefix):196 try:197 # Load the npz file198 data = np.load(npz_path)199 200 # Get the arrays and split them along the time dimension (T)201 poses = data["poses"]202 betas = data["betas"]203 expressions = data["expressions"]204 trans = data["trans"]205 206 # Determine the halfway point (T/2)207 half = poses.shape[0] // 2208 209 # Save the first half (0-5 seconds)210 np.savez(output_prefix + "_0_5.npz",211 betas=betas[:half],212 poses=poses[:half],213 expressions=expressions[:half],214 trans=trans[:half],215 model=data['model'],216 gender=data['gender'],217 mocap_frame_rate=data['mocap_frame_rate'])218 219 # Save the second half (5-10 seconds)220 np.savez(output_prefix + "_5_10.npz",221 betas=betas[half:],222 poses=poses[half:],223 expressions=expressions[half:],224 trans=trans[half:],225 model=data['model'],226 gender=data['gender'],227 mocap_frame_rate=data['mocap_frame_rate'])228 229 print(f"NPZ split saved for {output_prefix}")230 except Exception as e:231 print(f"Error processing NPZ file {npz_path}: {e}")232 233def split_wav(wav_path, output_prefix):234 try:235 with wave.open(wav_path, 'rb') as wav_file:236 params = wav_file.getparams()237 frames = wav_file.readframes(wav_file.getnframes())238 half_frame = len(frames) // 2239 240 # Create two half files241 for i, start_frame in enumerate([0, half_frame]):242 with wave.open(f"{output_prefix}_{i*5}_{(i+1)*5}.wav", 'wb') as out_wav:243 out_wav.setparams(params)244 if i == 0:245 out_wav.writeframes(frames[:half_frame])246 else:247 out_wav.writeframes(frames[half_frame:])248 print(f"WAV split saved for {output_prefix}")249 except Exception as e:250 print(f"Error processing WAV file {wav_path}: {e}")251 252def split_mp4(mp4_path, output_prefix):253 try:254 clip = VideoFileClip(mp4_path)255 for i in range(2):256 subclip = clip.subclip(i*5, (i+1)*5)257 subclip.write_videofile(f"{output_prefix}_{i*5}_{(i+1)*5}.mp4", codec="libx264", audio_codec="aac")258 print(f"MP4 split saved for {output_prefix}")259 except Exception as e:260 print(f"Error processing MP4 file {mp4_path}: {e}")261 262def process_files(root_dir, output_dir):263 import json264 clips = []265 dirs = os.listdir(root_dir)266 for dir in dirs:267 video_id = dir268 output_prefix = os.path.join(output_dir, video_id)269 root = os.path.join(root_dir, dir)270 npz_path = os.path.join(root, video_id + ".npz")271 wav_path = os.path.join(root, video_id + ".wav")272 mp4_path = os.path.join(root, video_id + ".mp4")273 274 # split_npz(npz_path, output_prefix)275 # split_wav(wav_path, output_prefix)276 # split_mp4(mp4_path, output_prefix)277 278 clip = {279 "video_id": video_id,280 "video_path": root,281 "audio_path": root,282 "motion_path": root,283 "mode": "test",284 "start_idx": 0,285 "end_idx": 150286 }287 clips.append(clip)288 289 output_json = output_dir + "/test.json"290 with open(output_json, 'w') as f:291 json.dump(clips, f, indent=4)292 293 294# Set the root directory path of your dataset and output directory295root_dir = '/content/oliver/oliver/Abortion_Laws_-_Last_Week_Tonight_with_John_Oliver_HBO-DRauXXz6t0Y.webm/test/'296output_dir = '/content/test'297 298# Make sure the output directory exists299os.makedirs(output_dir, exist_ok=True)300 301# Process all the files302process_files(root_dir, output_dir)303