| 220 | |
| 221 | class MBenchWiRefMotion(torch.utils.data.Dataset): |
| 222 | def __init__(self, test_json_file_list, motion_mean_path, motion_std_path, null_context_path, use_global_orient=True, text_key='prompt_video_detailed', test_seq_len=100, **kwargs): |
| 223 | base_json_file_list = test_json_file_list |
| 224 | self.data_list = [] |
| 225 | for json_file in base_json_file_list: |
| 226 | data_list = json.load(open(json_file, 'r')) |
| 227 | self.data_list.extend(data_list) |
| 228 | |
| 229 | print(len(self.data_list), "tensors cached in metadata.") |
| 230 | motion_mean = np.load(motion_mean_path) |
| 231 | motion_std = np.load(motion_std_path) |
| 232 | self.motion_mean = torch.from_numpy(motion_mean).float() |
| 233 | self.motion_std = torch.from_numpy(motion_std).float() |
| 234 | self.null_context = torch.load(null_context_path, weights_only=True, map_location="cpu") # [226, 4096] |
| 235 | self.motion_dim = motion_mean.shape[-1] |
| 236 | self.use_global_orient = use_global_orient |
| 237 | text_key_mapping = { |
| 238 | 'video_text_annot': 'prompt_video_detailed', |
| 239 | 'motion_text_annot': 'prompt_motion_detailed' |
| 240 | } |
| 241 | text_key = text_key_mapping.get(text_key, text_key) |
| 242 | self.text_key = text_key |
| 243 | self.prompt_emb_key = f'{text_key}_wanvideot5_embed_path' |
| 244 | self.test_seq_len = test_seq_len |
| 245 | |
| 246 | def __getitem__(self, index): |
| 247 | data = self.data_list[index] |