Spaces:
Sleeping
Sleeping
| import os | |
| import json | |
| import numpy as np | |
| import pandas as pd | |
| from tqdm import tqdm | |
| import librosa | |
| # Example parameters | |
| stride = 20 | |
| motion_length = 64 | |
| speaker_target = 2 | |
| use_additional = False | |
| root_dir = './beat_english_v2.0.0/' | |
| output_dir = "./datasets/data_json/" | |
| os.makedirs(output_dir, exist_ok=True) | |
| train_test_split_path = './beat_english_v2.0.0/train_test_split.csv' | |
| df = pd.read_csv(train_test_split_path) | |
| filtered_df = df[(df['id'].str.split('_').str[0].astype(int) == speaker_target) & (df['type'] != 'additional')] | |
| clips = [] | |
| for idx, row_item in tqdm(filtered_df.iterrows()): | |
| video_id = row_item['id'] | |
| mode = row_item['type'] | |
| # check exist | |
| npz_path = os.path.join(root_dir, "smplxflame_30", video_id + ".npz") | |
| wav_path = os.path.join(root_dir, "wave16k", video_id + ".wav") | |
| try: | |
| motion_data = np.load(npz_path, allow_pickle=True) | |
| except: | |
| print(f"cant open {npz_path}") | |
| try: | |
| wave_data, _ = librosa.load(wav_path, sr=None) | |
| except: | |
| print(f"cant open {wav_path}") | |
| motion = motion_data['poses'] | |
| total_len = motion.shape[0] | |
| for i in range(0, total_len - motion_length, stride): | |
| clip = { | |
| "video_id": video_id, | |
| "motion_path": npz_path, | |
| "audio_path": wav_path, | |
| "mode": mode, | |
| "start_idx": i, | |
| "end_idx": i + motion_length | |
| } | |
| clips.append(clip) | |
| output_json = os.path.join(output_dir, f"beat2_s{stride}_l{motion_length}_speaker{speaker_target}.json") | |
| with open(output_json, 'w') as f: | |
| json.dump(clips, f, indent=4) | |