| import h5py |
| import json |
| import numpy as np |
| from pathlib import Path |
| from sklearn.model_selection import StratifiedKFold, train_test_split |
|
|
| def generate_multitask_metadata(h5_input_path, json_output_path): |
| input_file = Path(h5_input_path) |
| output_file = Path(json_output_path) |
|
|
| if not input_file.exists(): |
| raise FileNotFoundError(f"Source data not found at: {input_file.resolve()}") |
|
|
| with h5py.File(input_file, 'r') as f: |
| emotion_labels = np.array(f['label_emotion']) |
| sign_labels = np.array(f['label_sign']) |
| |
| video_ids = np.arange(len(sign_labels)) |
|
|
| composite_labels = [f"{s}_{e}" for s, e in zip(sign_labels, emotion_labels)] |
|
|
| ids_train_val, ids_test, labels_train_val, labels_test = train_test_split( |
| video_ids, |
| composite_labels, |
| test_size=0.20, |
| random_state=42, |
| stratify=composite_labels |
| ) |
|
|
| skf = StratifiedKFold(n_splits=5, shuffle=True, random_state=42) |
| folds_data = [] |
|
|
| for train_idx, val_idx in skf.split(ids_train_val, labels_train_val): |
| folds_data.append({ |
| "train": [ |
| {"id": int(i), "sign": int(sign_labels[i]), "emotion": int(emotion_labels[i])} |
| for i in ids_train_val[train_idx] |
| ], |
| "val": [ |
| {"id": int(i), "sign": int(sign_labels[i]), "emotion": int(emotion_labels[i])} |
| for i in ids_train_val[val_idx] |
| ] |
| }) |
|
|
| metadata = { |
| "test_set": [ |
| {"id": int(i), "sign": int(sign_labels[i]), "emotion": int(emotion_labels[i])} |
| for i in ids_test |
| ], |
| "folds": folds_data |
| } |
|
|
| output_file.parent.mkdir(parents=True, exist_ok=True) |
|
|
| with open(output_file, 'w') as jf: |
| json.dump(metadata, jf, indent=4) |
| |
| print(f"Split metadata successfully written to: {output_file.resolve()}") |
|
|
| if __name__ == "__main__": |
| generate_multitask_metadata( |
| h5_input_path="fsl-data/multitask_mediapipe.h5", |
| json_output_path="metadata/multitask_splits.json" |
| ) |
|
|
