File size: 3,798 Bytes
bc971c7 | 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 | import os
import logging
import h5py
import json
import numpy as np
from tqdm import tqdm
from modules import mediapipe_generator
# --- Configuration ---
DATA_ROOT = "fsl-data/sliced_dataset"
HDF5_LOCATION = "fsl-data/multitask_mediapipe.h5"
SIGN_MAP_PATH = "metadata/sign_dict.json"
EMOTION_MAP_PATH = "metadata/emotion_dict.json"
logging.basicConfig(level=logging.INFO, format='%(asctime)s [%(levelname)s] %(message)s')
def get_mappings(root_path):
signs = sorted([d for d in os.listdir(root_path) if os.path.isdir(os.path.join(root_path, d))])
first_sign = os.path.join(root_path, signs[0])
emotions = sorted([d for d in os.listdir(first_sign) if os.path.isdir(os.path.join(first_sign, d))])
return {name: i for i, name in enumerate(signs)}, {name: i for i, name in enumerate(emotions)}
def save_json_mapping(mapping, filepath):
inverted_map = {int(v): k for k, v in mapping.items()}
with open(filepath, 'w', encoding='utf-8') as f:
json.dump(inverted_map, f, indent=4, ensure_ascii=False)
logging.info(f"Mapping saved to {filepath}")
def main():
sign_map, emotion_map = get_mappings(DATA_ROOT)
save_json_mapping(sign_map, SIGN_MAP_PATH)
save_json_mapping(emotion_map, EMOTION_MAP_PATH)
tasks = []
for sign in sign_map:
for emotion in emotion_map:
folder = os.path.join(DATA_ROOT, sign, emotion)
if os.path.exists(folder):
for video in os.listdir(folder):
tasks.append({
'path': os.path.join(folder, video),
'sign_id': sign_map[sign],
'emotion_id': emotion_map[emotion],
'name': video
})
num_samples = len(tasks)
str_dt = h5py.string_dtype(encoding='utf-8')
with h5py.File(HDF5_LOCATION, 'w') as f:
ds_l_sign = f.create_dataset('label_sign', (num_samples,), dtype='i')
ds_l_emotion = f.create_dataset('label_emotion', (num_samples,), dtype='i')
ds_l_path = f.create_dataset('label_filepath', (num_samples,), dtype=str_dt)
data_group = f.create_group('data')
for i, task in enumerate(tqdm(tasks, desc="Processing Videos")):
try:
data = mediapipe_generator.generate_mediapipe(filepath=task['path'])
if not data: continue
p_seq = np.array([mediapipe_generator.extract_to_array(r.pose_landmarks, 33, 4) for r in data])
f_seq = np.array([mediapipe_generator.extract_to_array(r.face_landmarks, 468, 3) for r in data])
lh_seq = np.array([mediapipe_generator.extract_to_array(r.left_hand_landmarks, 21, 3) for r in data])
rh_seq = np.array([mediapipe_generator.extract_to_array(r.right_hand_landmarks, 21, 3) for r in data])
sample_grp = data_group.create_group(str(i))
sample_grp.create_dataset('pose', data=p_seq, compression="gzip")
sample_grp.create_dataset('face', data=f_seq, compression="gzip")
sample_grp.create_dataset('left_hand', data=lh_seq, compression="gzip")
sample_grp.create_dataset('right_hand', data=rh_seq, compression="gzip")
ds_l_sign[i] = task['sign_id']
ds_l_emotion[i] = task['emotion_id']
ds_l_path[i] = task['path']
sample_grp.attrs['video_name'] = task['name']
sample_grp.attrs['frame_count'] = len(data)
except Exception as e:
logging.error(f"Error at {task['name']}: {e}")
f.attrs['sign_map_ref'] = SIGN_MAP_PATH
f.attrs['emotion_map_ref'] = EMOTION_MAP_PATH
if __name__ == "__main__":
main()
|