File size: 6,531 Bytes
721e031 |
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 |
{
"cells": [
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import os\n",
"import json\n",
"\n",
"\n",
"def collect_dataset_info(root_dir):\n",
" dataset_info = []\n",
"\n",
" for mode in [\"train\", \"val\", \"test\"]:\n",
" mode_dir = os.path.join(root_dir, mode)\n",
" if not os.path.exists(mode_dir):\n",
" continue\n",
"\n",
" for video_folder in os.listdir(mode_dir):\n",
" video_id = video_folder # Folder name is used as video_id\n",
" video_folder_path = os.path.join(mode_dir, video_folder)\n",
"\n",
" if os.path.isdir(video_folder_path):\n",
" video_path, audio_path, motion_path = None, None, None\n",
"\n",
" for file_name in os.listdir(video_folder_path):\n",
" file_path = os.path.join(video_folder_path, file_name)\n",
"\n",
" if file_name.endswith(\".mp4\"):\n",
" video_path = file_path\n",
" elif file_name.endswith(\".wav\"):\n",
" audio_path = file_path\n",
" elif file_name.endswith(\".pkl\"):\n",
" motion_path = file_path\n",
"\n",
" # Create an entry only if all the necessary files are present\n",
" if video_path and audio_path and motion_path:\n",
" dataset_info.append(\n",
" {\"video_id\": video_id, \"video_path\": video_path, \"audio_path\": audio_path, \"motion_path\": motion_path, \"mode\": mode}\n",
" )\n",
"\n",
" return dataset_info\n",
"\n",
"\n",
"# Set the root directory path of your dataset\n",
"root_dir = \"/path/to/ExpressiveWholeBodyDatasetReleaseV1.0\"\n",
"dataset_info = collect_dataset_info(root_dir)\n",
"output_file = \"dataset_info.json\"\n",
"with open(output_file, \"w\") as json_file:\n",
" json.dump(dataset_info, json_file, indent=4)\n",
"print(f\"Dataset information saved to {output_file}\")\n",
"\n",
"\n",
"import os\n",
"import json\n",
"import pickle\n",
"import wave\n",
"\n",
"\n",
"def load_pkl(pkl_path):\n",
" try:\n",
" with open(pkl_path, \"rb\") as f:\n",
" data = pickle.load(f)\n",
" return data\n",
" except Exception as e:\n",
" print(f\"Error loading {pkl_path}: {e}\")\n",
" return None\n",
"\n",
"\n",
"def load_wav(wav_path):\n",
" try:\n",
" with wave.open(wav_path, \"rb\") as f:\n",
" frames = f.getnframes()\n",
" return frames\n",
" except Exception as e:\n",
" print(f\"Error loading {wav_path}: {e}\")\n",
" return None\n",
"\n",
"\n",
"def generate_clips(data, stride, window_length):\n",
" clips = []\n",
" for entry in data:\n",
" pkl_data = load_pkl(entry[\"motion_path\"])\n",
" wav_frames = load_wav(entry[\"audio_path\"])\n",
"\n",
" # Only continue if both the pkl and wav files are successfully loaded\n",
" if pkl_data is None or wav_frames is None:\n",
" continue\n",
"\n",
" # Determine the total length of the sequence from pkl data\n",
" total_frames = len(pkl_data) # Assuming pkl contains motion data frames\n",
"\n",
" # Generate clips based on stride and window_length\n",
" for start_idx in range(0, total_frames - window_length + 1, stride):\n",
" end_idx = start_idx + window_length\n",
" clip = {\n",
" \"video_id\": entry[\"video_id\"],\n",
" \"video_path\": entry[\"video_path\"],\n",
" \"audio_path\": entry[\"audio_path\"],\n",
" \"motion_path\": entry[\"motion_path\"],\n",
" \"mode\": entry[\"mode\"],\n",
" \"start_idx\": start_idx,\n",
" \"end_idx\": end_idx,\n",
" }\n",
" clips.append(clip)\n",
"\n",
" return clips\n",
"\n",
"\n",
"# Load the existing dataset JSON file\n",
"input_json = \"dataset_info.json\"\n",
"with open(input_json, \"r\") as f:\n",
" dataset_info = json.load(f)\n",
"\n",
"# Set stride and window length\n",
"stride = 5 # Adjust stride as needed\n",
"window_length = 10 # Adjust window length as needed\n",
"\n",
"# Generate clips for all data\n",
"clips_data = generate_clips(dataset_info, stride, window_length)\n",
"\n",
"# Save the filtered clips data to a new JSON file\n",
"output_json = \"filtered_clips_data.json\"\n",
"with open(output_json, \"w\") as f:\n",
" json.dump(clips_data, f, indent=4)\n",
"\n",
"print(f\"Filtered clips data saved to {output_json}\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"import torch\n",
"import pickle\n",
"\n",
"with open(\n",
" \"/Users/haiyang/download_backup/oliver/Abortion_Laws_-_Last_Week_Tonight_with_John_Oliver_HBO-DRauXXz6t0Y.webm/test/214438-00_07_16-00_07_26/214438-00_07_16-00_07_26.pkl\",\n",
" \"rb\",\n",
") as f:\n",
" # Load the file by mapping any GPU tensors to CPU\n",
" pkl_example = torch.load(f, map_location=torch.device(\"cpu\"))\n",
"\n",
"# Now check the type of the object\n",
"print(type(pkl_example))\n",
"\n",
"# If it's a dictionary, print its keys\n",
"if isinstance(pkl_example, dict):\n",
" print(pkl_example.keys())"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.11.9"
}
},
"nbformat": 4,
"nbformat_minor": 2
}
|