dataset/dataset.py

import os
import random
import numpy as np
from decord import VideoReader
import glob
from tqdm import tqdm
import pickle
import torch
import torchvision.transforms as transforms
from torch.utils.data.dataset import Dataset
from decord import cpu, gpu
from torchvision.io import read_video, write_video

import json
import traceback

# MultiSample : 每个视频采样多个样本
# RandomRef : 参考图像是不是随机的，不是随机的就是上一帧,random ref的情况下，refvideo只有一帧
# MultiRef : 多帧参考帧，默认最多是8帧

class AMDConsecutiveVideo(Dataset):
    def __init__(
            self, 
            video_dir: str      = '', # video dir or pkl file
            sample_size: int    = 32, 
            sample_stride: int  = 2, 
            sample_n_frames:int = 16,
            ref_drop_ratio = 0.0,
        ):
        
        # Init setting
        self.sample_stride   = sample_stride
        self.sample_n_frames = sample_n_frames
        self.ref_drop_ratio = ref_drop_ratio
        
        # Transform
        sample_size = tuple(sample_size) if not isinstance(sample_size, int) else (sample_size, sample_size) # (256,256)
        self.pixel_transforms = transforms.Compose([
            transforms.Resize(sample_size[0]),
            transforms.CenterCrop(sample_size),
            transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True),
        ])

        if 'pkl' in video_dir:
            with open(video_dir, 'rb') as f:
                video_files = pickle.load(f)
            print(f'Total {len(video_files)} !!!')
        elif '.txt' in video_dir:
            with open(video_dir, 'r') as file:
                lines = file.readlines()
            video_dirs = [line.strip() for line in lines]

            video_files = []
            for dir in video_dirs:
                video_files += glob.glob(os.path.join(dir, '**', '*.mp4'), recursive=True)
            print(f'Total {len(video_files)} !!!')
    
        else:
            video_files = glob.glob(os.path.join(video_dir, '**', '*.mp4'), recursive=True)
            print(f'Total {len(video_files)} !!!')

        # Data dict
        self.metadata_list = []
        for file_path in tqdm(video_files):
            d = {}
            d['name'] = self.get_file_name(file_path)
            d['video_path'] = file_path
            self.metadata_list.append(d)
        self.length = len(self.metadata_list)
        print(f'Total {self.length} files is available')
    
    def __getitem__(self, idx):
        while True:
            try:
                file_name = self.metadata_list[idx]['name']
                file_name,videos,ref_img = self.get_batch(idx)
                break

            except Exception as e:
                # file_name = self.metadata_list[idx]['name']
                # print(file_name)
                print('error',e)
                idx = random.randint(0, self.length-1)

        sample = dict(name=file_name,videos=videos,ref_img = ref_img)
        return sample 

    def __len__(self):
        return self.length

    def get_batch(self, idx):
        
        # init
        meta_data = self.metadata_list[idx]
        file_name = meta_data['name']
        video_path = meta_data['video_path']

        # video process
        video_reader = VideoReader(video_path, ctx=cpu(0))
        video_length = len(video_reader)
        
        sample_frames = self.sample_n_frames + 1 # refimg + videos
        clip_length = min(video_length, (sample_frames - 1) * self.sample_stride + 1)
        start_idx   = random.randint(0, video_length - clip_length)
        batch_index = np.linspace(start_idx, start_idx + clip_length - 1, sample_frames, dtype=int)

        videos = torch.from_numpy(video_reader.get_batch(batch_index).asnumpy()).permute(0, 3, 1, 2).contiguous() #(N,H,W,C)->(N,C,H,W)
        videos = videos / 255.0 

        
        # transform
        videos_cache = self.pixel_transforms(videos) # F+1,C,H,W
        videos = videos_cache[1:,:,:,:] # F,C,H,W
        ref_frame = videos_cache[0,:,:,:] # C,H,W

        # repeat
        ref_frame = ref_frame.unsqueeze(0).repeat(videos.shape[0],1,1,1) # F,C,H,W

        
        return file_name,videos,ref_frame



    def get_file_name(self, file_path):
        return file_path.split('/')[-1].split('.')[0]
    
    @staticmethod
    def collate_fn(batch):
        # name
        name = [item['name'] for item in batch]

        # videos
        videos = [item['videos'] for item in batch]
        videos = torch.stack(videos)

        # ref_img
        ref_img = [item['ref_img'] for item in batch]
        ref_img = torch.stack(ref_img)

        randomref_img = None


        return dict(name=name, videos=videos, ref_img=ref_img,randomref_img=randomref_img)

class AMDConsecutiveVideoBalance(Dataset):
    def __init__(
            self, 
            video_dir: str      = '', # video dir or pkl file
            sample_size: int    = 32, 
            sample_stride: int  = 2, 
            sample_n_frames:int = 16,
            ref_drop_ratio = 0.0,
        ):
        
        # Init setting
        self.sample_stride   = sample_stride
        self.sample_n_frames = sample_n_frames
        self.ref_drop_ratio = ref_drop_ratio

        assert '.txt' in video_dir 

        with open(video_dir, 'r') as file:
            lines = file.readlines()
        video_paths = [line.strip() for line in lines]

        assert len(video_paths) == 2

        self.dataset1 = AMDConsecutiveVideo(video_dir = video_paths[0], # video dir or pkl file
                                            sample_size    = sample_size, 
                                            sample_stride  = sample_stride, 
                                            sample_n_frames = sample_n_frames,
                                            ref_drop_ratio = 0.0,)

        self.dataset2 = AMDConsecutiveVideo(video_dir = video_paths[1], # video dir or pkl file
                                            sample_size    = sample_size, 
                                            sample_stride  = sample_stride, 
                                            sample_n_frames = sample_n_frames,
                                            ref_drop_ratio = 0.0,)
        
        self.len1 = len(self.dataset1)
        self.len2 = len(self.dataset2)
        print(f'Total {self.len1 + self.len2} !!!')

    def __len__(self):
        # 设置为足够大的数值（例如两倍的最大数据集长度）
        return 2 * max(self.len1,self.len2)

    def __getitem__(self, idx):
        # 使用PyTorch的随机数生成器，保证多进程安全
        if torch.rand(1).item() < 0.5:
            # 从A中随机抽取样本
            a_idx = torch.randint(0, len(self.dataset1), (1,)).item()
            return self.dataset1[a_idx]
        else:
            # 从B中随机抽取样本
            b_idx = torch.randint(0, len(self.dataset2), (1,)).item()
            return self.dataset2[b_idx]

    @staticmethod
    def collate_fn(batch):
        # name
        name = [item['name'] for item in batch]

        # videos
        videos = [item['videos'] for item in batch]
        videos = torch.stack(videos)

        # ref_img
        ref_img = [item['ref_img'] for item in batch]
        ref_img = torch.stack(ref_img)

        randomref_img = None


        return dict(name=name, videos=videos, ref_img=ref_img,randomref_img=randomref_img)

class AMDConsecutiveVideoDoubleRef(Dataset):
    def __init__(
            self, 
            video_dir: str      = '', # video dir or pkl file
            sample_size: int    = 32, 
            sample_stride: int  = 2, 
            sample_n_frames:int = 16,
            ref_drop_ratio = 0.0,
        ):
        
        # Init setting
        self.sample_stride   = sample_stride
        self.sample_n_frames = sample_n_frames
        self.ref_drop_ratio = ref_drop_ratio
        
        # Transform
        sample_size = tuple(sample_size) if not isinstance(sample_size, int) else (sample_size, sample_size) # (256,256)
        self.pixel_transforms = transforms.Compose([
            transforms.Resize(sample_size[0]),
            transforms.CenterCrop(sample_size),
            transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True),
        ])

        if 'pkl' in video_dir:
            with open(video_dir, 'rb') as f:
                video_files = pickle.load(f)
            print(f'Total {len(video_files)} !!!')
        elif '.txt' in video_dir:
            with open(video_dir, 'r') as file:
                lines = file.readlines()
            video_dirs = [line.strip() for line in lines]

            video_files = []
            for dir in video_dirs:
                video_files += glob.glob(os.path.join(dir, '**', '*.mp4'), recursive=True)
            print(f'Total {len(video_files)} !!!')
    
        else:
            video_files = glob.glob(os.path.join(video_dir, '**', '*.mp4'), recursive=True)
            print(f'Total {len(video_files)} !!!')

        # Data dict
        self.metadata_list = []
        for file_path in tqdm(video_files):
            d = {}
            d['name'] = self.get_file_name(file_path)
            d['video_path'] = file_path
            self.metadata_list.append(d)
        self.length = len(self.metadata_list)
        print(f'Total {self.length} files is available')
    
    def __getitem__(self, idx):
        while True:
            try:
                file_name = self.metadata_list[idx]['name']
                file_name,videos,ref_img,randomref_img = self.get_batch(idx)
                break

            except Exception as e:
                # file_name = self.metadata_list[idx]['name']
                # print(file_name)
                print('error',e)
                idx = random.randint(0, self.length-1)

        sample = dict(name=file_name,videos=videos,ref_img = ref_img,randomref_img=randomref_img)
        return sample 

    def __len__(self):
        return self.length

    def get_batch(self, idx):
        
        # init
        meta_data = self.metadata_list[idx]
        file_name = meta_data['name']
        video_path = meta_data['video_path']

        # video process
        video_reader = VideoReader(video_path, ctx=cpu(0))
        video_length = len(video_reader)
        
        sample_frames = self.sample_n_frames + 1 # refimg + videos
        clip_length = min(video_length, (sample_frames - 1) * self.sample_stride + 1)
        start_idx   = random.randint(0, video_length - clip_length)
        batch_index = np.linspace(start_idx, start_idx + clip_length - 1, sample_frames, dtype=int)

        # random ref frame
        idx_all = np.arange(0,video_length)
        occ_idx = np.arange(start_idx, start_idx + clip_length)
        randomref_idx = [x for x in idx_all if x not in occ_idx]
        if len(randomref_idx) == 0:
            ref_frame_idx = batch_index[0]
        else:
            i = torch.randint(low=0, high=len(randomref_idx), size=(1,)).item()
            ref_frame_idx = randomref_idx[i]
        
        batch_index = [ref_frame_idx] + list(batch_index)

        videos = torch.from_numpy(video_reader.get_batch(batch_index).asnumpy()).permute(0, 3, 1, 2).contiguous() #(N,H,W,C)->(N,C,H,W)
        videos = videos / 255.0 

        # # ref frame
        # idx_all = np.arange(0,video_length)
        # occ_idx = np.arange(start_idx, start_idx + clip_length)
        # ref_idx = [x for x in idx_all if x not in occ_idx]
        # if len(ref_idx) == 0:
        #     ref_frame_idx = batch_index[0]
        # else:
        #     np.random.shuffle(ref_idx)
        #     ref_frame_idx = ref_idx[0]
        # ref_frame = torch.from_numpy(video_reader[ref_frame_idx].asnumpy()).permute(2, 0, 1).contiguous() 
        # ref_frame = ref_frame / 255.0
        
        # transform
        videos_cache = self.pixel_transforms(videos) # F+1,C,H,W
        videos = videos_cache[2:,:,:,:] # F,C,H,W
        ref_frame = videos_cache[1,:,:,:] # C,H,W
        randomref_frame = videos_cache[0,:,:,:] # C,H,W

        # repeat
        ref_frame = ref_frame.unsqueeze(0).repeat(videos.shape[0],1,1,1) # F,C,H,W
        randomref_frame = randomref_frame.unsqueeze(0).repeat(videos.shape[0],1,1,1) # F,C,H,W

        
        return file_name,videos,ref_frame,randomref_frame



    def get_file_name(self, file_path):
        return file_path.split('/')[-1].split('.')[0]
    
    @staticmethod
    def collate_fn(batch):
        # name
        name = [item['name'] for item in batch]

        # videos
        videos = [item['videos'] for item in batch]
        videos = torch.stack(videos)

        # ref_img
        ref_img = [item['ref_img'] for item in batch]
        ref_img = torch.stack(ref_img)

        randomref_img = [item['randomref_img'] for item in batch]
        randomref_img = torch.stack(randomref_img)


        return dict(name=name, videos=videos, ref_img=ref_img,randomref_img=randomref_img)

class AMDConsecutiveVideoDoubleRefBalance(Dataset):
    def __init__(
            self, 
            video_dir: str      = '', # video dir or pkl file
            sample_size: int    = 32, 
            sample_stride: int  = 2, 
            sample_n_frames:int = 16,
            ref_drop_ratio = 0.0,
        ):
        
        # Init setting
        self.sample_stride   = sample_stride
        self.sample_n_frames = sample_n_frames
        self.ref_drop_ratio = ref_drop_ratio

        assert '.txt' in video_dir 

        with open(video_dir, 'r') as file:
            lines = file.readlines()
        video_paths = [line.strip() for line in lines]

        self.datasets = []
        for vp in video_paths:
            self.datasets.append(AMDConsecutiveVideoDoubleRef(video_dir = vp, # video dir or pkl file
                                            sample_size    = sample_size, 
                                            sample_stride  = sample_stride, 
                                            sample_n_frames = sample_n_frames,
                                            ref_drop_ratio = 0.0,))
        
        self.length = len(self.datasets) * max([len(d) for d in self.datasets])

        print(self.length)

    def __len__(self):
        # 设置为足够大的数值（例如两倍的最大数据集长度）
        return self.length

    def __getitem__(self, idx):
        
        dataset_num = len(self.datasets)
        cur_idx = idx % dataset_num
        cur_dataset = self.datasets[cur_idx]
        idx = torch.randint(0, len(cur_dataset), (1,)).item()
        return cur_dataset[idx]

    @staticmethod
    def collate_fn(batch):
        # name
        name = [item['name'] for item in batch]

        # videos
        videos = [item['videos'] for item in batch]
        videos = torch.stack(videos)

        # ref_img
        ref_img = [item['ref_img'] for item in batch]
        ref_img = torch.stack(ref_img)
        
        randomref_img = [item['randomref_img'] for item in batch]
        randomref_img = torch.stack(randomref_img)


        return dict(name=name, videos=videos, ref_img=ref_img,randomref_img=randomref_img)



class AMDRandomPair(Dataset):
    def __init__(
            self, 
            video_dir: str      = '', # video dir or pkl file
            sample_size: int    = 32, 
            sample_stride: int  = 4, 
            sample_n_frames:int = 16,
            ref_drop_ratio = 0.0,
        ):
        
        # Init setting
        self.sample_stride   = sample_stride
        self.sample_n_frames = sample_n_frames
        self.ref_drop_ratio = ref_drop_ratio
        
        # Transform
        sample_size = tuple(sample_size) if not isinstance(sample_size, int) else (sample_size, sample_size) # (256,256)
        self.pixel_transforms = transforms.Compose([
            transforms.Resize(min(sample_size)),
            transforms.CenterCrop(sample_size),
            transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True),
        ])

        if 'pkl' in video_dir:
            with open(video_dir, 'rb') as f:
                video_files = pickle.load(f)
            print(f'Total {len(video_files)} !!!')
        elif '.txt' in video_dir:
            with open(video_dir, 'r') as file:
                lines = file.readlines()
            video_dirs = [line.strip() for line in lines]

            video_files = []
            for dir in video_dirs:
                video_files += glob.glob(os.path.join(dir, '**', '*.mp4'), recursive=True)
            print(f'Total {len(video_files)} !!!')
    
        else:
            video_files = glob.glob(os.path.join(video_dir, '**', '*.mp4'), recursive=True)

        # Data dict
        self.metadata_list = []
        for file_path in tqdm(video_files):
            d = {}
            d['name'] = self.get_file_name(file_path)
            d['video_path'] = file_path
            self.metadata_list.append(d)
        self.length = len(self.metadata_list)
        print(f'Total {self.length} files is available')
    
    def __getitem__(self, idx):
        while True:
            try:
                file_name = self.metadata_list[idx]['name']
                file_name,videos,ref_img = self.get_batch(idx)
                break

            except Exception as e:
                # file_name = self.metadata_list[idx]['name']
                # print(file_name)
                print('error',e)
                idx = random.randint(0, self.length-1)

        sample = dict(name=file_name,videos=videos,ref_img = ref_img)
        return sample 

    def __len__(self):
        return self.length

    def get_batch(self, idx):
        
        # init
        meta_data = self.metadata_list[idx]
        file_name = meta_data['name']
        video_path = meta_data['video_path']

        # video process
        video_reader = VideoReader(video_path)
        video_length = len(video_reader)
        

        ref_idx,video_idx = generate_non_equal_random_lists(frame_num=video_length,sample_num=self.sample_n_frames)

        ref_videos = torch.from_numpy(video_reader.get_batch(ref_idx).asnumpy()).permute(0, 3, 1, 2).contiguous() #(N,H,W,C)->(N,C,H,W)
        ref_videos = ref_videos / 255.0
        ref_videos = self.pixel_transforms(ref_videos)

        videos = torch.from_numpy(video_reader.get_batch(video_idx).asnumpy()).permute(0, 3, 1, 2).contiguous() #(N,H,W,C)->(N,C,H,W)
        videos = videos / 255.0
        videos = self.pixel_transforms(videos)

        return file_name,videos,ref_videos



    def get_file_name(self, file_path):
        return file_path.split('/')[-1].split('.')[0]
    
    @staticmethod
    def collate_fn(batch):
        # name
        name = [item['name'] for item in batch]

        # videos
        videos = [item['videos'] for item in batch]
        videos = torch.stack(videos)

        # ref_img
        ref_img = [item['ref_img'] for item in batch]
        ref_img = torch.stack(ref_img)

        return dict(name=name, videos=videos, ref_img=ref_img,randomref_img=None)

class A2MVideoAudio(Dataset):
    def __init__(
            self, 
            video_dir:str,
            sample_size: int    = 256, 
            sample_stride: int  = 1, 
            sample_n_frames:int = 16,
        ):
        super().__init__()

        self.sample_stride   = sample_stride
        self.sample_n_frames = sample_n_frames

        # Transform
        sample_size = tuple(sample_size) if not isinstance(sample_size, int) else (sample_size, sample_size) # (256,256)
        self.pixel_transforms = transforms.Compose([
            transforms.Resize(min(sample_size)),
            transforms.CenterCrop(sample_size),
            transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True),
        ])

        with open(video_dir, 'rb') as f:
            self.metadata_list = pickle.load(f)
        self.length = len(self.metadata_list)
        print(f'Total {self.length} files is available')
    
    def __getitem__(self, idx):
        while True:
            try:
                sample = self.get_batch(idx)
                break
            except Exception as e:
                print('error',e)
                idx = torch.randint(low=0, high=self.length, size=(1,)).item()

        return sample 

    def __len__(self):
        return self.length

    def get_batch(self, idx):

        """
        videos : 31,3,256,256
        ref_img : 3,256,256
        audio_feature : 30,50,384
        ref_pose : 3,256,256
        meta 
        """
        
        # init
        meta_data = self.metadata_list[idx]
        video_path = meta_data['video_path']
        whisper_path = meta_data['whisper_emb_path']

        # audio
        audio_feature = torch.load(whisper_path) 
        
        # load & check
        video_reader = VideoReader(video_path)
        video_length = min(len(video_reader),audio_feature.shape[0])
        
        # sample_frames
        sample_frames = self.sample_n_frames + 1 # self.sample_n_frames = 4, sample_frames=5
        clip_length = (sample_frames - 1) * self.sample_stride + 1 # clip_length = 9

        if clip_length > video_length :
            batch_index = np.linspace(0, clip_length - 1, sample_frames, dtype=int)
            batch_index = np.array([d for d in batch_index if d <= video_length-1],dtype=int)

            # frames
            videos = torch.from_numpy(video_reader.get_batch(batch_index).asnumpy()).permute(0, 3, 1, 2).contiguous() #(N,H,W,C)->(N,C,H,W)
            videos = videos / 255.0
            videos = self.pixel_transforms(videos)
            ref_video = videos[0,:] # C,H,W
            gt_video = videos[1:,:] # F,C,H,W

            audios = audio_feature[batch_index]
            ref_audio = audios[0,:]  # M,D
            gt_audio = audios[1:,:]  # F,M,D

            # available length
            cur_available_length =   gt_video.shape[0]

            # pad
            pad_length = self.sample_n_frames - gt_video.shape[0]
            video_pad = torch.zeros((pad_length, *gt_video.shape[1:]), dtype=gt_video.dtype)
            gt_video = torch.cat([gt_video, video_pad], dim=0) # F,C,H,W

            audio_pad = torch.zeros((pad_length, *gt_audio.shape[1:]), dtype=gt_audio.dtype)
            gt_audio = torch.cat([gt_audio, audio_pad], dim=0)
        else:
            start_idx = np.random.randint(0, video_length - clip_length + 1)
            end_idx = start_idx + clip_length
            batch_index = np.linspace(start_idx, end_idx - 1, sample_frames, dtype=int)

            # frames
            videos = torch.from_numpy(video_reader.get_batch(batch_index).asnumpy()).permute(0, 3, 1, 2).contiguous() #(N,H,W,C)->(N,C,H,W)
            videos = videos / 255.0
            videos = self.pixel_transforms(videos)
            ref_video = videos[0,:] # C,H,W
            gt_video = videos[1:,:] # F,C,H,W

            audios = audio_feature[batch_index]
            ref_audio = audios[0,:]  # M,D
            gt_audio = audios[1:,:]  # F,M,D

            # available length
            cur_available_length =   gt_video.shape[0]

        assert gt_video.shape[0] == self.sample_n_frames ,''+str(gt_video.shape[0])+' '+str(self.sample_n_frames)
        assert gt_audio.shape[0] == self.sample_n_frames ,''+str(gt_audio.shape[0])+' '+str(self.sample_n_frames)
        
        # mask
        mask = torch.zeros(self.sample_n_frames)
        mask[:cur_available_length] = 1
    
        # meta
        meta_ = dict(
            video_length = video_length,
            video_path = video_path,
            audio_path = whisper_path,
        )

        return dict(
            ref_video=ref_video,
            gt_video=gt_video,
            ref_audio=ref_audio,
            gt_audio=gt_audio,
            mask = mask,
            meta=meta_
        )
    
    def get_file_name(self, file_path):
        return file_path.split('/')[-1].split('.')[0]
    
    @staticmethod
    def collate_fn(batch):
        return dict(
            meta = [item["meta"] for item in batch],
            ref_video = torch.stack([item['ref_video'] for item in batch]),
            gt_video = torch.stack([item['gt_video'] for item in batch]),
            ref_audio = torch.stack([item['ref_audio'] for item in batch]),
            gt_audio = torch.stack([item['gt_audio'] for item in batch]),
            mask = torch.stack([item['mask'] for item in batch])
        )

class A2MVideoAudioPose(Dataset):
    def __init__(
            self, 
            video_dir:str,
            sample_size: int    = 256, 
            sample_stride: int  = 1, 
            sample_n_frames:int = 16,
            audio_drop_ratio:float = 0.0,
            **kwargs
        ):
        super().__init__()

        self.sample_stride   = sample_stride
        self.sample_n_frames = sample_n_frames
        self.audio_drop_ratio = audio_drop_ratio

        # Transform
        sample_size = tuple(sample_size) if not isinstance(sample_size, int) else (sample_size, sample_size) # (256,256)
        self.pixel_transforms = transforms.Compose([
            transforms.Resize(min(sample_size)),
            transforms.CenterCrop(sample_size),
            transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True),
        ])

        with open(video_dir, 'rb') as f:
            self.metadata_list = pickle.load(f)
        self.length = len(self.metadata_list)
        print(f'Total {self.length} files is available')
    
    def __getitem__(self, idx):
        while True:
            try:
                sample = self.get_batch(idx)
                break
            except Exception as e:
                print('error',e)
                idx = torch.randint(low=0, high=self.length, size=(1,)).item()

        return sample 

    def __len__(self):
        return self.length

    def get_batch(self, idx):

        """
        videos : 31,3,256,256
        ref_img : 3,256,256
        audio_feature : 30,50,384
        ref_pose : 3,256,256
        meta 
        """
        
        # init
        meta_data = self.metadata_list[idx]
        video_path = meta_data['video_path']
        whisper_path = meta_data['whisper_emb_path']
        pose_path = meta_data['pose_path']

        # audio
        audio_feature = torch.load(whisper_path) 
        
        # load & check
        video_reader = VideoReader(video_path)
        pose_reader = VideoReader(pose_path)
        video_length = min(len(video_reader),audio_feature.shape[0],len(pose_reader))
        
        # sample_frames
        sample_frames = self.sample_n_frames + 1 # self.sample_n_frames = 4, sample_frames=5
        clip_length = (sample_frames - 1) * self.sample_stride + 1 # clip_length = 9

        if clip_length > video_length :
            batch_index = np.linspace(0, clip_length - 1, sample_frames, dtype=int)
            batch_index = np.array([d for d in batch_index if d <= video_length-1],dtype=int)

            # frames
            videos = torch.from_numpy(video_reader.get_batch(batch_index).asnumpy()).permute(0, 3, 1, 2).contiguous() #(N,H,W,C)->(N,C,H,W)
            videos = videos / 255.0
            videos = self.pixel_transforms(videos)
            ref_video = videos[0,:] # C,H,W
            gt_video = videos[1:,:] # F,C,H,W

            poses = torch.from_numpy(pose_reader.get_batch(batch_index).asnumpy()).permute(0, 3, 1, 2).contiguous() #(N,H,W,C)->(N,C,H,W)
            poses = poses / 255.0
            poses = self.pixel_transforms(poses)
            ref_pose = poses[0,:] # C,H,W
            gt_pose = poses[1:,:] # F,C,H,W

            audios = audio_feature[batch_index]
            ref_audio = audios[0,:]  # M,D
            gt_audio = audios[1:,:]  # F,M,D

            # available length
            cur_available_length =   gt_video.shape[0]

            # pad
            pad_length = self.sample_n_frames - gt_video.shape[0]

            video_pad = torch.zeros((pad_length, *gt_video.shape[1:]), dtype=gt_video.dtype)
            gt_video = torch.cat([gt_video, video_pad], dim=0) # F,C,H,W

            pose_pad = torch.zeros((pad_length, *gt_pose.shape[1:]), dtype=gt_pose.dtype)
            gt_pose = torch.cat([gt_pose, pose_pad], dim=0) # F,C,H,W

            audio_pad = torch.zeros((pad_length, *gt_audio.shape[1:]), dtype=gt_audio.dtype)
            gt_audio = torch.cat([gt_audio, audio_pad], dim=0)
        else:
            start_idx = np.random.randint(0, video_length - clip_length + 1)
            end_idx = start_idx + clip_length
            batch_index = np.linspace(start_idx, end_idx - 1, sample_frames, dtype=int)

            # frames
            videos = torch.from_numpy(video_reader.get_batch(batch_index).asnumpy()).permute(0, 3, 1, 2).contiguous() #(N,H,W,C)->(N,C,H,W)
            videos = videos / 255.0
            videos = self.pixel_transforms(videos)
            ref_video = videos[0,:] # C,H,W
            gt_video = videos[1:,:] # F,C,H,W

            poses = torch.from_numpy(pose_reader.get_batch(batch_index).asnumpy()).permute(0, 3, 1, 2).contiguous() #(N,H,W,C)->(N,C,H,W)
            poses = poses / 255.0
            poses = self.pixel_transforms(poses)
            ref_pose = poses[0,:] # C,H,W
            gt_pose = poses[1:,:] # F,C,H,W            

            audios = audio_feature[batch_index]
            ref_audio = audios[0,:]  # M,D
            gt_audio = audios[1:,:]  # F,M,D

            # available length
            cur_available_length =   gt_video.shape[0]

        assert gt_video.shape[0] == self.sample_n_frames ,''+str(gt_video.shape[0])+' '+str(self.sample_n_frames)
        assert gt_audio.shape[0] == self.sample_n_frames ,''+str(gt_audio.shape[0])+' '+str(self.sample_n_frames)
        assert gt_pose.shape[0] == self.sample_n_frames ,''+str(gt_pose.shape[0])+' '+str(self.sample_n_frames)
        
        # mask
        mask = torch.zeros(self.sample_n_frames)
        mask[:cur_available_length] = 1

        # drop audio
        if torch.rand(1).item() < self.audio_drop_ratio:
            ref_audio = torch.zeros_like(ref_audio)
            gt_audio = torch.zeros_like(gt_audio)
    
        # meta
        meta_ = dict(
            video_length = video_length,
            video_path = video_path,
            audio_path = whisper_path,
        )

        return dict(
            ref_video=ref_video,
            gt_video=gt_video,
            ref_pose = ref_pose,
            gt_pose = gt_pose,
            ref_audio=ref_audio,
            gt_audio=gt_audio,
            mask = mask,
            meta=meta_
        )
    
    def get_file_name(self, file_path):
        return file_path.split('/')[-1].split('.')[0]
    
    @staticmethod
    def collate_fn(batch):
        return dict(
            meta = [item["meta"] for item in batch],
            ref_video = torch.stack([item['ref_video'] for item in batch]),
            gt_video = torch.stack([item['gt_video'] for item in batch]),
            ref_pose = torch.stack([item['ref_pose'] for item in batch]),
            gt_pose = torch.stack([item['gt_pose'] for item in batch]),
            ref_audio = torch.stack([item['ref_audio'] for item in batch]),
            gt_audio = torch.stack([item['gt_audio'] for item in batch]),
            mask = torch.stack([item['mask'] for item in batch])
        )

class A2MVideoAudioPoseRandomRef(Dataset):
    def __init__(
            self, 
            video_dir:str,
            sample_size: int    = 256, 
            sample_stride: int  = 1, 
            sample_n_frames:int = 16,
            **kwargs
        ):
        super().__init__()

        self.sample_stride   = sample_stride
        self.sample_n_frames = sample_n_frames

        # Transform
        sample_size = tuple(sample_size) if not isinstance(sample_size, int) else (sample_size, sample_size) # (256,256)
        self.pixel_transforms = transforms.Compose([
            transforms.Resize(min(sample_size)),
            transforms.CenterCrop(sample_size),
            transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True),
        ])

        with open(video_dir, 'rb') as f:
            self.metadata_list = pickle.load(f)
        self.length = len(self.metadata_list)
        print(f'Total {self.length} files is available')
    
    def __getitem__(self, idx):
        while True:
            try:
                sample = self.get_batch(idx)
                break
            except Exception as e:
                print('error',e)
                idx = torch.randint(low=0, high=self.length, size=(1,)).item()

        return sample 

    def __len__(self):
        return self.length

    def get_batch(self, idx):

        """
        videos : 31,3,256,256
        ref_img : 3,256,256
        audio_feature : 30,50,384
        ref_pose : 3,256,256
        meta 
        """
        
        # init
        meta_data = self.metadata_list[idx]
        video_path = meta_data['video_path']
        whisper_path = meta_data['whisper_emb_path']
        pose_path = meta_data['pose_path']

        # audio
        audio_feature = torch.load(whisper_path) 
        
        # load & check
        video_reader = VideoReader(video_path)
        pose_reader = VideoReader(pose_path)
        video_length = min(len(video_reader),audio_feature.shape[0],len(pose_reader))
        
        # sample_frames
        sample_frames = self.sample_n_frames 
        clip_length = (sample_frames - 1) * self.sample_stride + 1 # clip_length = 9

        if clip_length > video_length :
            batch_index = np.linspace(0, clip_length - 1, sample_frames, dtype=int)
            batch_index = list(np.array([d for d in batch_index if d <= video_length-1],dtype=int))

            # ref idx
            idx_all = np.arange(0,video_length)
            start_idx = 0
            occ_idx = np.arange(start_idx, start_idx + clip_length)
            ref_idx = [x for x in idx_all if x not in occ_idx]
            if len(ref_idx) == 0:
                ref_frame_idx = batch_index[0]
            else:
                np.random.shuffle(ref_idx)
                ref_frame_idx = ref_idx[0]
            batch_index = [ref_frame_idx] + batch_index

            # frames
            videos = torch.from_numpy(video_reader.get_batch(batch_index).asnumpy()).permute(0, 3, 1, 2).contiguous() #(N,H,W,C)->(N,C,H,W)
            videos = videos / 255.0
            videos = self.pixel_transforms(videos)
            ref_video = videos[0,:] # C,H,W
            gt_video = videos[1:,:] # F,C,H,W

            poses = torch.from_numpy(pose_reader.get_batch(batch_index).asnumpy()).permute(0, 3, 1, 2).contiguous() #(N,H,W,C)->(N,C,H,W)
            poses = poses / 255.0
            poses = self.pixel_transforms(poses)
            ref_pose = poses[0,:] # C,H,W
            gt_pose = poses[1:,:] # F,C,H,W

            audios = audio_feature[batch_index]
            ref_audio = audios[0,:]  # M,D
            gt_audio = audios[1:,:]  # F,M,D

            # available length
            cur_available_length =   gt_video.shape[0]

            # pad
            pad_length = self.sample_n_frames - gt_video.shape[0]

            video_pad = torch.zeros((pad_length, *gt_video.shape[1:]), dtype=gt_video.dtype)
            gt_video = torch.cat([gt_video, video_pad], dim=0) # F,C,H,W

            pose_pad = torch.zeros((pad_length, *gt_pose.shape[1:]), dtype=gt_pose.dtype)
            gt_pose = torch.cat([gt_pose, pose_pad], dim=0) # F,C,H,W

            audio_pad = torch.zeros((pad_length, *gt_audio.shape[1:]), dtype=gt_audio.dtype)
            gt_audio = torch.cat([gt_audio, audio_pad], dim=0)
        else:
            start_idx = np.random.randint(0, video_length - clip_length + 1)
            end_idx = start_idx + clip_length
            batch_index = list(np.linspace(start_idx, end_idx - 1, sample_frames, dtype=int))

            # ref index
            idx_all = np.arange(0,video_length)
            occ_idx = np.arange(start_idx, start_idx + clip_length)
            ref_idx = [x for x in idx_all if x not in occ_idx]
            if len(ref_idx) == 0:
                ref_frame_idx = batch_index[0]
            else:
                np.random.shuffle(ref_idx)
                ref_frame_idx = ref_idx[0]

            batch_index = [ref_frame_idx] + batch_index

            # frames
            videos = torch.from_numpy(video_reader.get_batch(batch_index).asnumpy()).permute(0, 3, 1, 2).contiguous() #(N,H,W,C)->(N,C,H,W)
            videos = videos / 255.0
            videos = self.pixel_transforms(videos)
            ref_video = videos[0,:] # C,H,W
            gt_video = videos[1:,:] # F,C,H,W

            poses = torch.from_numpy(pose_reader.get_batch(batch_index).asnumpy()).permute(0, 3, 1, 2).contiguous() #(N,H,W,C)->(N,C,H,W)
            poses = poses / 255.0
            poses = self.pixel_transforms(poses)
            ref_pose = poses[0,:] # C,H,W
            gt_pose = poses[1:,:] # F,C,H,W            

            audios = audio_feature[batch_index]
            ref_audio = audios[0,:]  # M,D
            gt_audio = audios[1:,:]  # F,M,D

            # available length
            cur_available_length =   gt_video.shape[0]

        assert gt_video.shape[0] == self.sample_n_frames ,''+str(gt_video.shape[0])+' '+str(self.sample_n_frames)
        assert gt_audio.shape[0] == self.sample_n_frames ,''+str(gt_audio.shape[0])+' '+str(self.sample_n_frames)
        assert gt_pose.shape[0] == self.sample_n_frames ,''+str(gt_pose.shape[0])+' '+str(self.sample_n_frames)
        
        # mask
        mask = torch.zeros(self.sample_n_frames)
        mask[:cur_available_length] = 1
    
        # meta
        meta_ = dict(
            video_length = video_length,
            video_path = video_path,
            audio_path = whisper_path,
        )

        return dict(
            ref_video=ref_video,
            gt_video=gt_video,
            ref_pose = ref_pose,
            gt_pose = gt_pose,
            ref_audio=ref_audio,
            gt_audio=gt_audio,
            mask = mask,
            meta=meta_
        )
    
    def get_file_name(self, file_path):
        return file_path.split('/')[-1].split('.')[0]
    
    @staticmethod
    def collate_fn(batch):
        return dict(
            meta = [item["meta"] for item in batch],
            ref_video = torch.stack([item['ref_video'] for item in batch]),
            gt_video = torch.stack([item['gt_video'] for item in batch]),
            ref_pose = torch.stack([item['ref_pose'] for item in batch]),
            gt_pose = torch.stack([item['gt_pose'] for item in batch]),
            ref_audio = torch.stack([item['ref_audio'] for item in batch]),
            gt_audio = torch.stack([item['gt_audio'] for item in batch]),
            mask = torch.stack([item['mask'] for item in batch])
        )

class A2MVideoAudioPoseMultiSample(Dataset):
    def __init__(
            self, 
            video_dir:str,
            sample_size: int    = 256, 
            sample_stride: int  = 1, 
            sample_n_frames:int = 16,
            audio_drop_ratio:float = 0.0,
            num_sample:int = 4,
            **kwargs
        ):
        super().__init__()

        self.sample_stride   = sample_stride
        self.sample_n_frames = sample_n_frames
        self.audio_drop_ratio = audio_drop_ratio
        self.num_sample = num_sample

        # Transform
        sample_size = tuple(sample_size) if not isinstance(sample_size, int) else (sample_size, sample_size) # (256,256)
        self.pixel_transforms = transforms.Compose([
            transforms.Resize(min(sample_size)),
            transforms.CenterCrop(sample_size),
            transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True),
        ])

        with open(video_dir, 'rb') as f:
            self.metadata_list = pickle.load(f)

        self.length =len(self.metadata_list)
        print(f'Total {self.length} files is available')
    
    def __getitem__(self, idx):
        while True:
            try:

                sample = self.get_batch(idx)
                break
            except Exception as e:
                print('error',e)
                idx = torch.randint(low=0, high=self.length, size=(1,)).item()

        return sample 

    def __len__(self):
        return self.length

    def get_batch(self, idx):

        """
        videos : 31,3,256,256
        ref_img : 3,256,256
        audio_feature : 30,50,384
        ref_pose : 3,256,256
        meta 
        """
        
        # init
        meta_data = self.metadata_list[idx]
        video_path = meta_data['video_path']
        whisper_path = meta_data['whisper_emb_path']
        pose_path = meta_data['pose_path']

        # audio
        audio_feature = torch.load(whisper_path) 
        
        # load & check
        video_reader = VideoReader(video_path)
        pose_reader = VideoReader(pose_path)
        video_length = min(len(video_reader),audio_feature.shape[0],len(pose_reader))
        
        # sample_frames
        sample_frames = self.sample_n_frames + 1 # self.sample_n_frames = 4, sample_frames=5
        clip_length = (sample_frames - 1) * self.sample_stride + 1 # clip_length = 9

        ref_video_list = []
        gt_video_list = []
        ref_pose_list = []
        gt_pose_list = []
        ref_audio_list = []
        gt_audio_list = []
        mask_list = []

        for i in range(self.num_sample):
            start_idx = np.random.randint(0, video_length - clip_length + 1)
            end_idx = start_idx + clip_length
            batch_index = np.linspace(start_idx, end_idx - 1, sample_frames, dtype=int)

            # frames
            videos = torch.from_numpy(video_reader.get_batch(batch_index).asnumpy()).permute(0, 3, 1, 2).contiguous() #(N,H,W,C)->(N,C,H,W)
            videos = videos / 255.0
            videos = self.pixel_transforms(videos)
            ref_video = videos[0,:] # C,H,W
            gt_video = videos[1:,:] # F,C,H,W

            poses = torch.from_numpy(pose_reader.get_batch(batch_index).asnumpy()).permute(0, 3, 1, 2).contiguous() #(N,H,W,C)->(N,C,H,W)
            poses = poses / 255.0
            poses = self.pixel_transforms(poses)
            ref_pose = poses[0,:] # C,H,W
            gt_pose = poses[1:,:] # F,C,H,W            

            audios = audio_feature[batch_index]
            ref_audio = audios[0,:]  # M,D
            gt_audio = audios[1:,:]  # F,M,D

            # available length
            cur_available_length =   gt_video.shape[0]

            assert gt_video.shape[0] == self.sample_n_frames ,''+str(gt_video.shape[0])+' '+str(self.sample_n_frames)
            assert gt_audio.shape[0] == self.sample_n_frames ,''+str(gt_audio.shape[0])+' '+str(self.sample_n_frames)
            assert gt_pose.shape[0] == self.sample_n_frames ,''+str(gt_pose.shape[0])+' '+str(self.sample_n_frames)
            
            # mask
            mask = torch.zeros(self.sample_n_frames)
            mask[:cur_available_length] = 1

            # drop audio
            if torch.rand(1).item() < self.audio_drop_ratio:
                ref_audio = torch.zeros_like(ref_audio)
                gt_audio = torch.zeros_like(gt_audio)

            # cache
            ref_video_list.append(ref_video)
            gt_video_list.append(gt_video)
            ref_pose_list.append(ref_pose)
            gt_pose_list.append(gt_pose)
            ref_audio_list.append(ref_audio)
            gt_audio_list.append(gt_audio)
            mask_list.append(mask)
    

        return dict(
            ref_video=torch.stack(ref_video_list),
            gt_video=torch.stack(gt_video_list),
            ref_pose = torch.stack(ref_pose_list),
            gt_pose = torch.stack(gt_pose_list),
            ref_audio=torch.stack(ref_audio_list),
            gt_audio=torch.stack(gt_audio_list),
            mask = torch.stack(mask_list),
        )
    
    def get_file_name(self, file_path):
        return file_path.split('/')[-1].split('.')[0]
    
    @staticmethod
    def collate_fn(batch):
        return dict(
            ref_video = torch.cat([item['ref_video'] for item in batch],dim=0),
            gt_video = torch.cat([item['gt_video'] for item in batch],dim=0),
            ref_pose = torch.cat([item['ref_pose'] for item in batch],dim=0),
            gt_pose = torch.cat([item['gt_pose'] for item in batch],dim=0),
            ref_audio = torch.cat([item['ref_audio'] for item in batch],dim=0),
            gt_audio = torch.cat([item['gt_audio'] for item in batch],dim=0),
            mask = torch.cat([item['mask'] for item in batch],dim=0)
        )


class A2MVideoAudioPoseMultiSampleMultiRefBalance(Dataset):
    def __init__(
            self, 
            video_dir:str,
            sample_size: int    = 256, 
            sample_stride: int  = 1, 
            sample_n_frames:int = 16,
            audio_drop_ratio:float = 0.0,
            num_sample:int = 4,
            max_ref_frame:int = 8,
            random_ref_num:bool = False,
            **kwargs
        ):
        super().__init__()


        with open(video_dir, 'r') as file:
            lines = file.readlines()
        video_dirs = [line.strip() for line in lines]
        assert len(video_dirs) == 2, 'Only support 2 video dirs'

        self.dataset1 = A2MVideoAudioPoseMultiSampleMultiRef(video_dir=video_dirs[0],
                                                     sample_size=sample_size,
                                                     sample_stride=sample_stride,
                                                     sample_n_frames=sample_n_frames,
                                                     audio_drop_ratio=audio_drop_ratio,
                                                     num_sample=num_sample,
                                                     max_ref_frame=max_ref_frame,
                                                     random_ref_num=random_ref_num)
        self.dataset2 = A2MVideoAudioPoseMultiSampleMultiRef(video_dir=video_dirs[1],
                                                     sample_size=sample_size,
                                                     sample_stride=sample_stride,
                                                     sample_n_frames=sample_n_frames,
                                                     audio_drop_ratio=audio_drop_ratio,
                                                     num_sample=num_sample,
                                                     max_ref_frame=max_ref_frame,
                                                     random_ref_num=random_ref_num)
        self.length = 2*max(len(self.dataset1),len(self.dataset2))

    
    def __getitem__(self, idx):
        while True:
            try:
                if idx % 2 == 0:
                    a_idx = torch.randint(0, len(self.dataset1), (1,)).item()
                    sample = self.dataset1[a_idx]
                else:
                    b_idx = torch.randint(0, len(self.dataset2), (1,)).item()
                    sample = self.dataset2[b_idx]
                break
            except Exception as e:
                print('error',e)
                idx = torch.randint(low=0, high=self.length, size=(1,)).item()

        return sample 

    def __len__(self):
        return self.length

    @staticmethod
    def collate_fn(batch):
        return dict(
            ref_video = torch.cat([item['ref_video'] for item in batch],dim=0),
            gt_video = torch.cat([item['gt_video'] for item in batch],dim=0),
            ref_pose = torch.cat([item['ref_pose'] for item in batch],dim=0),
            gt_pose = torch.cat([item['gt_pose'] for item in batch],dim=0),
            ref_audio = torch.cat([item['ref_audio'] for item in batch],dim=0),
            gt_audio = torch.cat([item['gt_audio'] for item in batch],dim=0),
            mask = torch.cat([item['mask'] for item in batch],dim=0)
        )

class A2MVideoAudioMultiRefDoubleRef(Dataset):
    def __init__(
            self, 
            video_dir:str,
            sample_size: int    = 256, 
            sample_stride: int  = 1, 
            sample_n_frames:int = 16,
            audio_drop_ratio:float = 0.0,
            num_sample:int = 4,
            max_ref_frame:int = 8,
            **kwargs
        ):
        super().__init__()

        self.sample_stride   = sample_stride
        self.sample_n_frames = sample_n_frames
        self.audio_drop_ratio = audio_drop_ratio
        self.num_sample = num_sample
        self.max_ref_frame = max_ref_frame
        self.randomref_num = 8

        # Transform
        sample_size = tuple(sample_size) if not isinstance(sample_size, int) else (sample_size, sample_size) # (256,256)
        self.pixel_transforms = transforms.Compose([
            transforms.Resize(min(sample_size)),
            transforms.CenterCrop(sample_size),
            transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True),
        ])


        with open(video_dir, 'rb') as f:
            self.metadata_list = pickle.load(f)
        self.length = len(self.metadata_list)
        print(f'Total {self.length} files is available')
    
    def __getitem__(self, idx):
        while True:
            try:
                sample = self.get_batch(idx)
                break
            except Exception as e:
                print('error',e)
                idx = torch.randint(low=0, high=self.length, size=(1,)).item()
                

        return sample 

    def __len__(self):
        return self.length

    def get_batch(self, idx):

        """
        videos : 31,3,256,256
        ref_img : 3,256,256
        audio_feature : 30,50,384
        ref_pose : 3,256,256
        meta 
        """
        
        # init
        meta_data = self.metadata_list[idx]
        video_path = meta_data['video_path']
        whisper_path = meta_data['whisper_emb_path']

        # audio
        audio_feature = torch.load(whisper_path) 
        
        # load & check
        video_reader = VideoReader(video_path)
        video_length = min(len(video_reader),audio_feature.shape[0])
        
        # ref num
        r = torch.rand(1).item()
        if r < 0.33:
            ref_num = 0
        if r < 0.66:
            ref_num = 1
        else:
            ref_num = self.max_ref_frame

        sample_frames = self.sample_n_frames + ref_num 
        clip_length = (sample_frames - 1) * self.sample_stride + 1 

        start_idx = np.random.randint(0, video_length - clip_length + 1)
        end_idx = start_idx + clip_length
        batch_index = list(np.linspace(start_idx, end_idx - 1, sample_frames, dtype=int))

        # randomref
        random_index = list(np.linspace(0, video_length - 1, self.randomref_num, dtype=int))

        # random_index = torch.randint(low=0, high=clip_length-1, size=(1,)).item()
        video_batch_index = random_index + batch_index

        # frames
        videos = torch.from_numpy(video_reader.get_batch(video_batch_index).asnumpy()).permute(0, 3, 1, 2).contiguous() #(N,H,W,C)->(N,C,H,W)
        videos = videos / 255.0
        videos = self.pixel_transforms(videos)

        randomref_video = videos[:self.randomref_num,:] # T,C,H,W
        l_videos = videos[self.randomref_num:,:] # T,C,H,W
        ref_video = l_videos[:ref_num,:] if ref_num > 0 else None
        gt_video = l_videos[ref_num:,:] # F,C,H,W
        
        audios = audio_feature[batch_index]
        ref_audio = audios[:ref_num,:]  if ref_num > 0 else None
        gt_audio = audios[ref_num:,:]  # F,M,D

        # padding ref frame
        if ref_num == 1:
            ref_video_pad = torch.zeros((self.max_ref_frame-ref_num,*ref_video.shape[1:]))
            ref_video = torch.cat([ref_video_pad,ref_video],dim=0) # N,T,C,H,W

            ref_audio_pad = torch.zeros((self.max_ref_frame-ref_num,*ref_audio.shape[1:]))
            ref_audio = torch.cat([ref_audio_pad,ref_audio],dim=0)
        elif ref_num == 0:
            ref_video = torch.zeros((self.max_ref_frame,*gt_video.shape[1:]))
            ref_audio = torch.zeros((self.max_ref_frame,*gt_audio.shape[1:]))




        # available length
        cur_available_length =   gt_video.shape[0]

        assert gt_video.shape[0] == self.sample_n_frames ,''+str(gt_video.shape[0])+' '+str(self.sample_n_frames)
        assert gt_audio.shape[0] == self.sample_n_frames ,''+str(gt_audio.shape[0])+' '+str(self.sample_n_frames)
        
        # mask
        mask = torch.zeros(self.sample_n_frames)
        mask[:cur_available_length] = 1



        return dict(
            ref_video=ref_video,
            gt_video=gt_video,
            randomref_video = randomref_video,
            ref_audio= ref_audio,
            gt_audio=gt_audio,
            mask = mask,
        )
    
    def get_file_name(self, file_path):
        return file_path.split('/')[-1].split('.')[0]
    
    @staticmethod
    def collate_fn(batch):
        return dict(
            ref_video = torch.stack([item['ref_video'] for item in batch],dim=0),
            gt_video = torch.stack([item['gt_video'] for item in batch],dim=0),
            randomref_video = torch.stack([item['randomref_video'] for item in batch],dim=0),
            ref_audio = torch.stack([item['ref_audio'] for item in batch],dim=0),
            gt_audio = torch.stack([item['gt_audio'] for item in batch],dim=0),
            mask = torch.stack([item['mask'] for item in batch],dim=0)
        )


class A2MVideoAudioMultiRefDoubleRefBalance(Dataset):
    def __init__(
            self, 
            video_dir:str,
            sample_size: int    = 256, 
            sample_stride: int  = 1, 
            sample_n_frames:int = 16,
            audio_drop_ratio:float = 0.0,
            num_sample:int = 4,
            max_ref_frame:int = 8,
            **kwargs
        ):
        super().__init__()

        self.sample_stride   = sample_stride
        self.sample_n_frames = sample_n_frames
        self.audio_drop_ratio = audio_drop_ratio
        self.num_sample = num_sample
        self.max_ref_frame = max_ref_frame
        self.randomref_num = 8

        # Transform
        sample_size = tuple(sample_size) if not isinstance(sample_size, int) else (sample_size, sample_size) # (256,256)
        self.pixel_transforms = transforms.Compose([
            transforms.Resize(min(sample_size)),
            transforms.CenterCrop(sample_size),
            transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True),
        ])

        with open(video_dir, 'rb') as f:
            self.metadata_list = pickle.load(f)
        self.length = len(self.metadata_list)
        print(f'Total {self.length} files is available')
    
    def __getitem__(self, idx):
        while True:
            try:
                sample = self.get_batch(idx)
                break
            except Exception as e:
                print('error',e)
                idx = torch.randint(low=0, high=self.length, size=(1,)).item()
                

        return sample 

    def __len__(self):
        return self.length

    def get_batch(self, idx):

        """
        videos : 31,3,256,256
        ref_img : 3,256,256
        audio_feature : 30,50,384
        ref_pose : 3,256,256
        meta 
        """
        
        # init
        meta_data = self.metadata_list[idx]
        video_path = meta_data['video_path']
        whisper_path = meta_data['whisper_emb_path']

        # audio
        audio_feature = torch.load(whisper_path) 
        
        # load & check
        video_reader = VideoReader(video_path)
        video_length = min(len(video_reader),audio_feature.shape[0])
        
        # ref num
        r = torch.rand(1).item()
        if r < 0.33:
            ref_num = 0
        elif r<0.66:
            ref_num = 1
        else:
            ref_num = self.max_ref_frame

        sample_frames = self.sample_n_frames + ref_num 
        clip_length = (sample_frames - 1) * self.sample_stride + 1 

        start_idx = np.random.randint(0, video_length - clip_length + 1)
        end_idx = start_idx + clip_length
        batch_index = list(np.linspace(start_idx, end_idx - 1, sample_frames, dtype=int))

        # randomref
        random_index = list(np.linspace(0, video_length - 1, self.randomref_num, dtype=int))
        random_index = list(random_index)

        video_batch_index = random_index + batch_index

        # frames
        videos = torch.from_numpy(video_reader.get_batch(video_batch_index).asnumpy()).permute(0, 3, 1, 2).contiguous() #(N,H,W,C)->(N,C,H,W)
        videos = videos / 255.0
        videos = self.pixel_transforms(videos)
        randomref_video = videos[:self.randomref_num,:] # T,C,H,W
        l_videos = videos[self.randomref_num:,:] # T,C,H,W
        ref_video = l_videos[:ref_num,:] if ref_num > 0 else None
        gt_video = l_videos[ref_num:,:] # F,C,H,W
        
        audios = audio_feature[batch_index]
        ref_audio = audios[:ref_num,:]  if ref_num > 0 else None
        gt_audio = audios[ref_num:,:]  # F,M,D

        # padding ref frame
        if ref_num == 1:
            ref_video_pad = torch.zeros((self.max_ref_frame-ref_num,*ref_video.shape[1:]))
            ref_video = torch.cat([ref_video_pad,ref_video],dim=0) # N,T,C,H,W

            ref_audio_pad = torch.zeros((self.max_ref_frame-ref_num,*ref_audio.shape[1:]))
            ref_audio = torch.cat([ref_audio_pad,ref_audio],dim=0)
        elif ref_num == 0:
            ref_video = torch.zeros((self.max_ref_frame,*gt_video.shape[1:]))
            ref_audio = torch.zeros((self.max_ref_frame,*gt_audio.shape[1:]))


        # available length
        cur_available_length =   gt_video.shape[0]

        assert gt_video.shape[0] == self.sample_n_frames ,''+str(gt_video.shape[0])+' '+str(self.sample_n_frames)
        assert gt_audio.shape[0] == self.sample_n_frames ,''+str(gt_audio.shape[0])+' '+str(self.sample_n_frames)
        
        # mask
        mask = torch.zeros(self.sample_n_frames)
        mask[:cur_available_length] = 1



        return dict(
            ref_video=ref_video,
            gt_video=gt_video,
            randomref_video = randomref_video,
            ref_audio= ref_audio,
            gt_audio=gt_audio,
            mask = mask,
        )
    
    def get_file_name(self, file_path):
        return file_path.split('/')[-1].split('.')[0]
    
    @staticmethod
    def collate_fn(batch):
        return dict(
            ref_video = torch.stack([item['ref_video'] for item in batch],dim=0),
            gt_video = torch.stack([item['gt_video'] for item in batch],dim=0),
            randomref_video = torch.stack([item['randomref_video'] for item in batch],dim=0),
            ref_audio = torch.stack([item['ref_audio'] for item in batch],dim=0),
            gt_audio = torch.stack([item['gt_audio'] for item in batch],dim=0),
            mask = torch.stack([item['mask'] for item in batch],dim=0)
        )



# pose img2img
class A2MVideoAudioPoseRandomRefMultiSample(Dataset):
    def __init__(
            self, 
            video_dir:str,
            sample_size: int    = 256, 
            sample_stride: int  = 1, 
            sample_n_frames:int = 16,
            num_sample:int = 4,
            max_ref_frame:int = 8,
            random_ref_num:bool = False,
            **kwargs
        ):
        super().__init__()

        self.sample_stride   = sample_stride
        self.sample_n_frames = sample_n_frames
        self.num_sample = num_sample
        self.max_ref_frame = max_ref_frame

        # Transform
        sample_size = tuple(sample_size) if not isinstance(sample_size, int) else (sample_size, sample_size) # (256,256)
        self.pixel_transforms = transforms.Compose([
            transforms.Resize(min(sample_size)),
            transforms.CenterCrop(sample_size),
            transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True),
        ])

        with open(video_dir, 'rb') as f:
            self.metadata_list = pickle.load(f)

        self.length = len(self.metadata_list)
        print(f'Total {self.length} files is available')
    
    def __getitem__(self, idx):
        while True:
            try:
                sample = self.get_batch(idx)
                break
            except Exception as e:
                print('error',e)
                idx = torch.randint(low=0, high=self.length, size=(1,)).item()

        return sample 

    def __len__(self):
        return self.length

    def get_batch(self, idx):

        """
        videos : 31,3,256,256
        ref_img : 3,256,256
        audio_feature : 30,50,384
        ref_pose : 3,256,256
        meta 
        """
        
        # init
        meta_data = self.metadata_list[idx]
        video_path = meta_data['video_path']
        whisper_path = meta_data['whisper_emb_path']
        pose_path = meta_data['pose_path']

        # audio
        audio_feature = torch.load(whisper_path) 
        
        # load & check
        video_reader = VideoReader(video_path)
        pose_reader = VideoReader(pose_path)
        video_length = min(len(video_reader),audio_feature.shape[0],len(pose_reader))
        
        # batch index
        sample_frames = self.num_sample * 2
        if video_length < sample_frames:
            raise ValueError(f"视频长度{video_length}太短了,需要长度{sample_frames}")
        # occ_idx = np.arange(0, video_length)
        # # # np.random.shuffle(occ_idx)
        # # # 生成随机排列的索引
        # # shuffled_indices = torch.randperm(len(occ_idx))
        # # # 使用索引打乱列表
        # # occ_idx = [occ_idx[i] for i in shuffled_indices]

        # occ_idx = shuffle_list(occ_idx)
        # batch_index = occ_idx[:sample_frames]

        # batch_index = [torch.randint(low=0, high=video_length, size=(1,)).item() for i in range(sample_frames)]
        start_idx = torch.randint(low=0, high=video_length-sample_frames, size=(1,)).item()
        occ_idx = np.arange(start_idx, start_idx+video_length)
        batch_index = occ_idx[:sample_frames]


        # frames
        videos = torch.from_numpy(video_reader.get_batch(batch_index).asnumpy()).permute(0, 3, 1, 2).contiguous() #(N,H,W,C)->(N,C,H,W)
        videos = videos / 255.0
        videos = self.pixel_transforms(videos)
        ref_video = videos[:self.num_sample,:].unsqueeze(1) # N,1,C,H,W
        gt_video = videos[self.num_sample:,:].unsqueeze(1) #  N,1,C,H,W

        poses = torch.from_numpy(pose_reader.get_batch(batch_index).asnumpy()).permute(0, 3, 1, 2).contiguous() #(N,H,W,C)->(N,C,H,W)
        poses = poses / 255.0
        poses = self.pixel_transforms(poses)
        ref_pose = poses[:self.num_sample,:].unsqueeze(1) #  N,1,C,H,W
        gt_pose = poses[self.num_sample:,:].unsqueeze(1) #  N,1,C,H,W 

        audios = audio_feature[batch_index]
        ref_audio = audios[:self.num_sample,:].unsqueeze(1)  #  N,1,C,H,W
        gt_audio = audios[self.num_sample:,:].unsqueeze(1)  #  N,1,C,H,W
          

        # available length
        cur_available_length =   gt_video.shape[0]
        
        # mask
        mask = torch.zeros(self.sample_n_frames)
        mask[:cur_available_length] = 1

        return dict(
            ref_video=ref_video,
            gt_video=gt_video,
            ref_pose = ref_pose,
            gt_pose =gt_pose,
            ref_audio=ref_audio,
            gt_audio=gt_audio,
            mask = mask,
        )

    
    def get_file_name(self, file_path):
        return file_path.split('/')[-1].split('.')[0]
    
    @staticmethod
    def collate_fn(batch):
        return dict(
            ref_video = torch.cat([item['ref_video'] for item in batch],dim=0),
            gt_video = torch.cat([item['gt_video'] for item in batch],dim=0),
            ref_pose = torch.cat([item['ref_pose'] for item in batch],dim=0),
            gt_pose = torch.cat([item['gt_pose'] for item in batch],dim=0),
            ref_audio = torch.cat([item['ref_audio'] for item in batch],dim=0),
            gt_audio = torch.cat([item['gt_audio'] for item in batch],dim=0),
            mask = torch.cat([item['mask'] for item in batch],dim=0)
        )


class A2MVideoAudioPoseMultiSampleMultiRef(Dataset):
    def __init__(
            self, 
            video_dir:str,
            sample_size: int    = 256, 
            sample_stride: int  = 1, 
            sample_n_frames:int = 16,
            audio_drop_ratio:float = 0.0,
            num_sample:int = 4,
            max_ref_frame:int = 8,
            random_ref_num:bool = False,
            **kwargs
        ):
        super().__init__()

        self.sample_stride   = sample_stride
        self.sample_n_frames = sample_n_frames
        self.audio_drop_ratio = audio_drop_ratio
        self.num_sample = num_sample
        self.max_ref_frame = max_ref_frame
        self.random_ref_num = random_ref_num

        # Transform
        sample_size = tuple(sample_size) if not isinstance(sample_size, int) else (sample_size, sample_size) # (256,256)
        self.pixel_transforms = transforms.Compose([
            transforms.Resize(min(sample_size)),
            transforms.CenterCrop(sample_size),
            transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True),
        ])

        with open(video_dir, 'rb') as f:
            self.metadata_list = pickle.load(f)
        self.length = len(self.metadata_list)
        print(f'Total {self.length} files is available')
    
    def __getitem__(self, idx):
        while True:
            try:
                sample = self.get_batch(idx)
                break
            except Exception as e:
                print('error',e)
                idx = torch.randint(low=0, high=self.length, size=(1,)).item()
                

        return sample 

    def __len__(self):
        return self.length

    def get_batch(self, idx):

        """
        videos : 31,3,256,256
        ref_img : 3,256,256
        audio_feature : 30,50,384
        ref_pose : 3,256,256
        meta 
        """
        
        # init
        meta_data = self.metadata_list[idx]
        video_path = meta_data['video_path']
        whisper_path = meta_data['whisper_emb_path']
        pose_path = meta_data['pose_path']

        # audio
        audio_feature = torch.load(whisper_path) 
        
        # load & check
        video_reader = VideoReader(video_path)
        pose_reader = VideoReader(pose_path)
        video_length = min(len(video_reader),audio_feature.shape[0],len(pose_reader))
        
        # sample_frames


        ref_video_list = []
        gt_video_list = []
        ref_pose_list = []
        gt_pose_list = []
        ref_audio_list = []
        gt_audio_list = []
        mask_list = []

        for i in range(self.num_sample):
            # random ref num
            if self.random_ref_num:
                ref_num = torch.randint(low=1, high=self.max_ref_frame+1, size=(1,)).item()
            else:
                ref_num = [1, self.max_ref_frame][torch.randint(2, (1,)).item()]
            sample_frames = self.sample_n_frames + ref_num 
            clip_length = (sample_frames - 1) * self.sample_stride + 1 

            start_idx = np.random.randint(0, video_length - clip_length + 1)
            end_idx = start_idx + clip_length
            batch_index = np.linspace(start_idx, end_idx - 1, sample_frames, dtype=int)

            # frames
            videos = torch.from_numpy(video_reader.get_batch(batch_index).asnumpy()).permute(0, 3, 1, 2).contiguous() #(N,H,W,C)->(N,C,H,W)
            videos = videos / 255.0
            videos = self.pixel_transforms(videos)
            ref_video = videos[:ref_num,:] # T,C,H,W
            gt_video = videos[ref_num:,:] # F,C,H,W

            poses = torch.from_numpy(pose_reader.get_batch(batch_index).asnumpy()).permute(0, 3, 1, 2).contiguous() #(N,H,W,C)->(N,C,H,W)
            poses = poses / 255.0
            poses = self.pixel_transforms(poses)
            ref_pose = poses[:ref_num,:] # T,C,H,W
            gt_pose = poses[ref_num:,:] # F,C,H,W            

            audios = audio_feature[batch_index]
            ref_audio = audios[:ref_num,:]  # T,M,D
            gt_audio = audios[ref_num:,:]  # F,M,D

            # padding ref frame
            if ref_num < self.max_ref_frame:
                ref_video_pad = torch.zeros((self.max_ref_frame-ref_num,*ref_video.shape[1:]))
                ref_video = torch.cat([ref_video_pad,ref_video],dim=0) # N,T,C,H,W

                ref_pose_pad = torch.zeros((self.max_ref_frame-ref_num,*ref_pose.shape[1:]))
                ref_pose = torch.cat([ref_pose_pad,ref_pose],dim=0)

                ref_audio_pad = torch.zeros((self.max_ref_frame-ref_num,*ref_audio.shape[1:]))
                ref_audio = torch.cat([ref_audio_pad,ref_audio],dim=0)

            # available length
            cur_available_length =   gt_video.shape[0]

            assert gt_video.shape[0] == self.sample_n_frames ,''+str(gt_video.shape[0])+' '+str(self.sample_n_frames)
            assert gt_audio.shape[0] == self.sample_n_frames ,''+str(gt_audio.shape[0])+' '+str(self.sample_n_frames)
            assert gt_pose.shape[0] == self.sample_n_frames ,''+str(gt_pose.shape[0])+' '+str(self.sample_n_frames)
            
            # mask
            mask = torch.zeros(self.sample_n_frames)
            mask[:cur_available_length] = 1

            # drop audio
            if torch.rand(1).item() < self.audio_drop_ratio:
                ref_audio = torch.zeros_like(ref_audio)
                gt_audio = torch.zeros_like(gt_audio)

            # cache
            ref_video_list.append(ref_video)
            gt_video_list.append(gt_video)
            ref_pose_list.append(ref_pose)
            gt_pose_list.append(gt_pose)
            ref_audio_list.append(ref_audio)
            gt_audio_list.append(gt_audio)
            mask_list.append(mask)
    

        return dict(
            ref_video=torch.stack(ref_video_list),
            gt_video=torch.stack(gt_video_list),
            ref_pose = torch.stack(ref_pose_list),
            gt_pose = torch.stack(gt_pose_list),
            ref_audio=torch.stack(ref_audio_list),
            gt_audio=torch.stack(gt_audio_list),
            mask = torch.stack(mask_list),
        )
    
    def get_file_name(self, file_path):
        return file_path.split('/')[-1].split('.')[0]
    
    @staticmethod
    def collate_fn(batch):
        return dict(
            ref_video = torch.cat([item['ref_video'] for item in batch],dim=0),
            gt_video = torch.cat([item['gt_video'] for item in batch],dim=0),
            ref_pose = torch.cat([item['ref_pose'] for item in batch],dim=0),
            gt_pose = torch.cat([item['gt_pose'] for item in batch],dim=0),
            ref_audio = torch.cat([item['ref_audio'] for item in batch],dim=0),
            gt_audio = torch.cat([item['gt_audio'] for item in batch],dim=0),
            mask = torch.cat([item['mask'] for item in batch],dim=0)
        )

# inference
class A2VDataset(Dataset):
    def __init__(
            self, 
            video_dir:str,
            sample_size: int    = 256, 
            sample_stride: int  = 1, 
            sample_n_frames:int = 120,
            **kwargs
        ):
        super().__init__()

        self.sample_stride   = sample_stride
        self.sample_n_frames = sample_n_frames

        # Transform
        sample_size = tuple(sample_size) if not isinstance(sample_size, int) else (sample_size, sample_size) # (256,256)
        self.pixel_transforms = transforms.Compose([
            transforms.Resize(min(sample_size)),
            transforms.CenterCrop(sample_size),
            transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True),
        ])

        with open(video_dir, 'rb') as f:
            self.metadata_list = pickle.load(f)
        self.length = len(self.metadata_list)
        print(f'Total {self.length} files is available')
    
    def __getitem__(self, idx):
        while True:
            try:
                sample = self.get_batch(idx)
                break
            except Exception as e:
                print('error',e)
                idx = torch.randint(low=0, high=self.length, size=(1,)).item()
                

        return sample 

    def __len__(self):
        return self.length

    def get_batch(self, idx):
        
        # init
        meta_data = self.metadata_list[idx]
        video_path = meta_data['video_path']
        whisper_path = meta_data['whisper_emb_path']
        audio_path = meta_data['audio_path']
        # pose_path = meta_data['pose_path']
        name = os.path.basename(video_path).split('.')[0]
        fps = 25 if 'hdtf' in video_path else 30

        # audio
        audio_feature = torch.load(whisper_path) 
        
        # load & check
        video_reader = VideoReader(video_path)
        # pose_reader = VideoReader(pose_path)
        video_length = min(len(video_reader),audio_feature.shape[0])
        
        # sample_frame

        # batch idx
        ref_num = 1
        sample_frames = self.sample_n_frames + ref_num
        clip_length = (sample_frames - 1) * self.sample_stride + 1 

        start_idx = np.random.randint(0, video_length - clip_length + 1)
        end_idx = start_idx + clip_length
        batch_index = np.linspace(start_idx, end_idx - 1, sample_frames, dtype=int)

        # frames
        videos = torch.from_numpy(video_reader.get_batch(batch_index).asnumpy()).permute(0, 3, 1, 2).contiguous() #(N,H,W,C)->(N,C,H,W)
        videos = videos / 255.0
        videos = self.pixel_transforms(videos)
        ref_img = videos[:ref_num,:] # 1,C,H,W
        inf_video = videos[ref_num:,:] # F,C,H,W

        # poses = torch.from_numpy(pose_reader.get_batch(batch_index).asnumpy()).permute(0, 3, 1, 2).contiguous() #(N,H,W,C)->(N,C,H,W)
        # poses = poses / 255.0
        # poses = self.pixel_transforms(poses)
        # ref_pose = poses[:ref_num,:] # T,C,H,W
        # gt_pose = poses[ref_num:,:] # F,C,H,W            

        audios = audio_feature[batch_index]
        ref_audio = audios[:ref_num,:]  # 1,M,D
        inf_audio = audios[ref_num:,:]  # F,M,D

        # meta_info
        start_time = start_idx / fps
        _meta = {"name":name,
                 "audio_path":audio_path,
                 "fps":fps,
                 "start_time":start_time}

        return dict(
            meta_info = _meta,
            ref_img=ref_img,
            gt_video=videos,
            ref_audio=ref_audio,
            inf_audio=inf_audio,
            ref_pose=None,
            inf_pose=None,
        )
    
    @staticmethod
    def collate_fn(batch):
        return dict(
            meta_info = [item['meta_info'] for item in batch],
            ref_img= torch.stack([item["ref_img"] for item in batch]),
            ref_audio= torch.stack([item["ref_audio"] for item in batch]),
            inf_audio= torch.stack([item["inf_audio"] for item in batch]),
            ref_pose= torch.ones((2,2)),
            inf_pose= torch.ones((2,2)),
            gt_video= torch.stack([item["gt_video"] for item in batch]),
        )


def generate_non_equal_random_lists(frame_num,sample_num):
    list1 = [np.random.randint(0, frame_num) for _ in range(sample_num)]

    list2 = []
    for i in range(len(list1)):
        available_numbers = list(range(0, list1[i])) + list(range(list1[i] + 1, frame_num))
        list2.append(random.choice(available_numbers))

    return list1, list2

def shuffle_list(l):
    shuffled_indices = torch.randperm(len(l))
    # 使用索引打乱列表
    l = [l[i] for i in shuffled_indices]

    return l

if __name__ == "__main__":
    # dataset = CelebvText()
    
    
    # dataloader = torch.utils.data.DataLoader(dataset, batch_size=4, num_workers=16,)
    # for idx, batch in enumerate(dataloader):
        # print(batch["videos"].shape, len(batch["text"]))
        # for i in range(batch["videos"].shape[0]):
            # save_videos_grid(batch["videos"][i:i+1].permute(0,2,1,3,4), os.path.join(".", f"{idx}-{i}.mp4"), rescale=True)

    from torch.utils.data import DataLoader
    # dataset = AMDVideoAudioFeature(
    #     path=data_path,
    #     path_type="file",
    #     motion_seq_len=motion_seq_len,
    #     sample_n_frames=num_frames,
    #     audio_processor=audio_processor
    # )
    dataset = AMDVideoAudioFeature(
        video_dir = '/mnt/pfs-mc0p4k/tts/team/digital_avatar_group/sunwenzhang/qiyuan/code/AMD_linear/dataset/path/lhz/train.pkl',
        path_type = 'file'
    )
    dataloader = DataLoader(
        dataset,2,True,num_workers=0,
        collate_fn=dataset.collate_fn
    )

    # dataloader = DataLoader(dataset,batch_size=2,collate_fn=dataset.collate_fn,num_workers=2)
    
    # d = dataset[10]
    # video = d["videos"]
    # audio = d["audio_feature"]
    # refimg = d["ref_img"]
    for d in dataloader:
        video = d["videos"]
        audio_feature = d["audio_feature"]
        refimg = d["ref_img"]
        # break
    print(video.shape)
    print(audio_feature.shape)




#