name

fit/tools/load.py

@@ -1,25 +1,46 @@
 import scipy.io
 import numpy as np
-import json
+import json  # 引入json模块
 
 
 def load(name, path):
+    # 处理UTD-MHAD的JSON文件（你的单帧数据）
     if name == 'UTD_MHAD':
-        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
+        # 判断文件是否为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] # down_sample
+        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}")