| import os |
| import logging |
| import h5py |
| import json |
| import numpy as np |
| from tqdm import tqdm |
| from modules import mediapipe_generator |
|
|
| |
| 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() |
|
|
