File size: 2,141 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 | 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 correspond to the index in the label datasets
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"
)
|