46 lines
2.0 KiB
Python
46 lines
2.0 KiB
Python
import scipy.io
|
||
import numpy as np
|
||
import json # 引入json模块
|
||
|
||
|
||
def load(name, path):
|
||
# 处理UTD-MHAD的JSON文件（你的单帧数据）
|
||
if name == 'UTD_MHAD':
|
||
# 判断文件是否为JSON格式（通过后缀）
|
||
if path.endswith('.json'):
|
||
with open(path, 'r') as f:
|
||
data = json.load(f) # 加载JSON列表
|
||
# 转换为NumPy数组，确保形状为[关节数, 3]
|
||
data_np = np.array(data)
|
||
# 校验数据格式（防止错误）
|
||
assert data_np.ndim == 2 and data_np.shape[1] == 3, \
|
||
f"UTD-MHAD JSON格式错误，应为[关节数, 3]，实际为{data_np.shape}"
|
||
# 若需要单帧维度（[1, 关节数, 3]），可扩展维度
|
||
return data_np[np.newaxis, ...] # 输出形状：[1, N, 3]（1表示单帧）
|
||
|
||
# 保留原UTD_MHAD的.mat文件支持（如果还需要处理.mat数据）
|
||
elif path.endswith('.mat'):
|
||
arr = scipy.io.loadmat(path)['d_skel']
|
||
new_arr = np.zeros([arr.shape[2], arr.shape[0], arr.shape[1]])
|
||
for i in range(arr.shape[2]):
|
||
for j in range(arr.shape[0]):
|
||
for k in range(arr.shape[1]):
|
||
new_arr[i][j][k] = arr[j][k][i]
|
||
return new_arr
|
||
|
||
else:
|
||
raise ValueError(f"UTD-MHAD不支持的文件格式：{path}")
|
||
|
||
# 其他数据集的原有逻辑保持不变
|
||
elif name == 'HumanAct12':
|
||
return np.load(path, allow_pickle=True)
|
||
elif name == "CMU_Mocap":
|
||
return np.load(path, allow_pickle=True)
|
||
elif name == "Human3.6M":
|
||
return np.load(path, allow_pickle=True)[0::5] # 下采样
|
||
elif name == "NTU":
|
||
return np.load(path, allow_pickle=True)[0::2]
|
||
elif name == "HAA4D":
|
||
return np.load(path, allow_pickle=True)
|
||
else:
|
||
raise ValueError(f"不支持的数据集名称：{name}") |