Class MBenchWiRefMotion

datasets/video_datasets.py:221–319

class MBenchWiRefMotion(torch.utils.data.Dataset):
    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):
        base_json_file_list = test_json_file_list
        self.data_list = []
        for json_file in base_json_file_list:
            data_list = json.load(open(json_file, 'r'))
            self.data_list.extend(data_list)

        print(len(self.data_list), "tensors cached in metadata.")
        motion_mean = np.load(motion_mean_path)
        motion_std = np.load(motion_std_path)
        self.motion_mean = torch.from_numpy(motion_mean).float()
        self.motion_std = torch.from_numpy(motion_std).float()
        self.null_context = torch.load(null_context_path, weights_only=True, map_location="cpu")   # [226, 4096]
        self.motion_dim = motion_mean.shape[-1]
        self.use_global_orient = use_global_orient
        text_key_mapping = {
            'video_text_annot': 'prompt_video_detailed',
            'motion_text_annot': 'prompt_motion_detailed'
        }
        text_key = text_key_mapping.get(text_key, text_key)
        self.text_key = text_key
        self.prompt_emb_key = f'{text_key}_wanvideot5_embed_path'
        self.test_seq_len = test_seq_len
        
    def __getitem__(self, index):
        data = self.data_list[index]
        enable_m2m = False
        if "motion_path" in data and data.get("use_ref_motion", False):
            motion_path = data["motion_path"]
            motion = torch.load(motion_path, weights_only=True, map_location="cpu")
            if isinstance(motion, dict):
                motion = motion["motion"]
            # interpolate motion from 16 fps to 20 fps
            original_frames = motion.shape[0]
            target_frames = original_frames * 20 // 16
            # [original_frames, motion_dim] -> [target_frames, motion_dim]
            motion = torch.nn.functional.interpolate(motion.unsqueeze(0).permute(0, 2, 1), size=target_frames, mode='linear', align_corners=True).squeeze(0).permute(1, 0)
            enable_m2m = True
        else:
            motion = torch.zeros((self.test_seq_len, self.motion_dim)).float()
            
        motion_length = motion.shape[0]
        motion_duration = motion_length / 20.0  # assuming 20 fps
        
        # normalize motion
        data_dict = {}
        motion = motion[:, :self.motion_dim]
        motion = (motion - self.motion_mean) / self.motion_std
        data_dict["motion"] = motion
        data_dict["motion_length"] = motion_length
        motion_dim_mask = torch.ones(self.motion_dim)

        joint_num = 22
        # ref motion
        ref_motion = motion.clone()

        # only use the local ref motion