ldm/data/video_spec_maa2_dataset.py

import csv
import os
import pickle
import sys

import numpy as np
import torch
import random
import math
import librosa

class audio_video_spec_fullset_Dataset(torch.utils.data.Dataset):
    # Only Load audio dataset: for training Stage1: Audio Npy Dataset
    def __init__(self, split, dataset1, feat_type='clip', transforms=None, sr=22050, duration=10, truncate=220000, fps=21.5, drop=0.0, fix_frames=False, hop_len=256):
        super().__init__()

        if split == "train":
            self.split = "Train"

        elif split == "valid" or split == 'test':
            self.split = "Test"

        # Default params:
        self.min_duration = 2
        self.sr = sr                # 22050
        self.duration = duration    # 10
        self.truncate = truncate    # 220000
        self.fps = fps
        self.fix_frames = fix_frames
        self.hop_len = hop_len
        self.drop = drop
        print("Fix Frames: {}".format(self.fix_frames))
        print("Use Drop: {}".format(self.drop))

        # Dataset1: (VGGSound)
        assert dataset1.dataset_name == "VGGSound"
        
        # spec_dir: spectrogram path
        # feat_dir: CAVP feature path
        # video_dir: video path
        
        dataset1_spec_dir = os.path.join(dataset1.data_dir, "mel_maa2", "npy")
        dataset1_feat_dir = os.path.join(dataset1.data_dir, "cavp")
        dataset1_video_dir = os.path.join(dataset1.video_dir, "tmp_vid")
        
        split_txt_path = dataset1.split_txt_path
        with open(os.path.join(split_txt_path, '{}.txt'.format(self.split)), "r") as f:
            data_list1 = f.readlines()
            data_list1 = list(map(lambda x: x.strip(), data_list1))

            spec_list1 = list(map(lambda x: os.path.join(dataset1_spec_dir, x) + "_mel.npy", data_list1))      # spec
            feat_list1 = list(map(lambda x: os.path.join(dataset1_feat_dir, x) + ".npz",     data_list1))      # feat
            video_list1 = list(map(lambda x: os.path.join(dataset1_video_dir, x) + "_new_fps_21.5_truncate_0_10.0.mp4",   data_list1))      # video


        # Merge Data:
        self.data_list = data_list1 if self.split != "Test" else data_list1[:200]
        self.spec_list = spec_list1 if self.split != "Test" else spec_list1[:200]
        self.feat_list = feat_list1 if self.split != "Test" else feat_list1[:200]
        self.video_list = video_list1 if self.split != "Test" else video_list1[:200]

        assert len(self.data_list) == len(self.spec_list) == len(self.feat_list) == len(self.video_list)
        
        shuffle_idx = np.random.permutation(np.arange(len(self.data_list)))
        self.data_list = [self.data_list[i] for i in shuffle_idx]
        self.spec_list = [self.spec_list[i] for i in shuffle_idx]
        self.feat_list = [self.feat_list[i] for i in shuffle_idx]
        self.video_list = [self.video_list[i] for i in shuffle_idx]

        print('Split: {}  Sample Num: {}'.format(split, len(self.data_list)))



    def __len__(self):
        return len(self.data_list)
    

    def load_spec_and_feat(self, spec_path, video_feat_path):
        """Load audio spec and video feat"""
        try:
            spec_raw = np.load(spec_path).astype(np.float32)                    # channel: 1
        except:
            print(f"corrupted mel: {spec_path}", flush=True)
            spec_raw = np.zeros((80, 625), dtype=np.float32) # [C, T]

        p = np.random.uniform(0,1)
        if p > self.drop:
            try:
                video_feat = np.load(video_feat_path)['feat'].astype(np.float32)
            except:
                print(f"corrupted video: {video_feat_path}", flush=True)
                video_feat = np.load(os.path.join(os.path.dirname(video_feat_path), 'empty_vid.npz'))['feat'].astype(np.float32)
        else:
            video_feat = np.load(os.path.join(os.path.dirname(video_feat_path), 'empty_vid.npz'))['feat'].astype(np.float32)

        spec_len = self.sr * self.duration / self.hop_len
        if spec_raw.shape[1] < spec_len:
            spec_raw = np.tile(spec_raw, math.ceil(spec_len / spec_raw.shape[1]))
        spec_raw = spec_raw[:, :int(spec_len)]
        
        feat_len = self.fps * self.duration
        if video_feat.shape[0] < feat_len:
            video_feat = np.tile(video_feat, (math.ceil(feat_len / video_feat.shape[0]), 1))
        video_feat = video_feat[:int(feat_len)]
        return spec_raw, video_feat


    def mix_audio_and_feat(self, spec1=None, spec2=None, video_feat1=None, video_feat2=None, video_info_dict={}, mode='single'):
        """ Return Mix Spec and Mix video feat"""
        if mode == "single":
            # spec1:
            if not self.fix_frames:
                start_idx = random.randint(0, self.sr * self.duration - self.truncate - 1)  # audio start
            else:
                start_idx = 0

            start_frame = int(self.fps * start_idx / self.sr)
            truncate_frame = int(self.fps * self.truncate / self.sr)

            # Spec Start & Truncate:
            spec_start = int(start_idx / self.hop_len)
            spec_truncate = int(self.truncate / self.hop_len)

            spec1 = spec1[:, spec_start : spec_start + spec_truncate]
            video_feat1 = video_feat1[start_frame: start_frame + truncate_frame]

            # info_dict:
            video_info_dict['video_time1'] = str(start_frame) + '_' + str(start_frame+truncate_frame)   # Start frame, end frame
            video_info_dict['video_time2'] = ""
            return spec1, video_feat1, video_info_dict
        
        elif mode == "concat":
            total_spec_len = int(self.truncate / self.hop_len)
            # Random Trucate len:
            spec1_truncate_len = random.randint(self.min_duration * self.sr // self.hop_len, total_spec_len - self.min_duration * self.sr // self.hop_len - 1)
            spec2_truncate_len = total_spec_len - spec1_truncate_len

            # Sample spec clip:
            spec_start1 = random.randint(0, total_spec_len - spec1_truncate_len - 1)
            spec_start2 = random.randint(0, total_spec_len - spec2_truncate_len - 1)
            spec_end1, spec_end2 = spec_start1 + spec1_truncate_len, spec_start2 + spec2_truncate_len
            
            # concat spec:
            spec1, spec2 = spec1[:, spec_start1 : spec_end1], spec2[:, spec_start2 : spec_end2]
            concat_audio_spec = np.concatenate([spec1, spec2], axis=1)  

            # Concat Video Feat:
            start1_frame, truncate1_frame = int(self.fps * spec_start1 * self.hop_len / self.sr), int(self.fps * spec1_truncate_len * self.hop_len / self.sr)
            start2_frame, truncate2_frame = int(self.fps * spec_start2 * self.hop_len / self.sr), int(self.fps * self.truncate / self.sr) - truncate1_frame
            video_feat1, video_feat2 = video_feat1[start1_frame : start1_frame + truncate1_frame], video_feat2[start2_frame : start2_frame + truncate2_frame]
            concat_video_feat = np.concatenate([video_feat1, video_feat2])

            video_info_dict['video_time1'] = str(start1_frame) + '_' + str(start1_frame+truncate1_frame)   # Start frame, end frame
            video_info_dict['video_time2'] = str(start2_frame) + '_' + str(start2_frame+truncate2_frame)
            return concat_audio_spec, concat_video_feat, video_info_dict 



    def __getitem__(self, idx):
        
        audio_name1 = self.data_list[idx]
        spec_npy_path1 = self.spec_list[idx]
        video_feat_path1 = self.feat_list[idx]
        video_path1 = self.video_list[idx]

        # select other video:
        flag = False
        if random.uniform(0, 1) < 0.5:
            flag = True
            random_idx = idx
            while random_idx == idx:
                random_idx = random.randint(0, len(self.data_list)-1)
            audio_name2 = self.data_list[random_idx]
            spec_npy_path2 = self.spec_list[random_idx]
            video_feat_path2 = self.feat_list[random_idx]
            video_path2 = self.video_list[random_idx]

        # Load the Spec and Feat:
        spec1, video_feat1 = self.load_spec_and_feat(spec_npy_path1, video_feat_path1)

        if flag:
            spec2, video_feat2 = self.load_spec_and_feat(spec_npy_path2, video_feat_path2)
            video_info_dict = {'audio_name1':audio_name1, 'audio_name2': audio_name2, 'video_path1': video_path1, 'video_path2': video_path2}
            mix_spec, mix_video_feat, mix_info = self.mix_audio_and_feat(spec1, spec2, video_feat1, video_feat2, video_info_dict, mode='concat')
        else:
            video_info_dict = {'audio_name1':audio_name1, 'audio_name2': "", 'video_path1': video_path1, 'video_path2': ""}
            mix_spec, mix_video_feat, mix_info = self.mix_audio_and_feat(spec1=spec1, video_feat1=video_feat1, video_info_dict=video_info_dict, mode='single')

        # print("mix spec shape:", mix_spec.shape)
        # print("mix video feat:", mix_video_feat.shape)
        data_dict = {}
        # data_dict['mix_spec'] = mix_spec[None].repeat(3, axis=0) # TODO：要把这里改掉，否则无法适应maa的autoencoder
        data_dict['mix_spec'] = mix_spec # (80, 512)
        data_dict['mix_video_feat'] = mix_video_feat # (32, 512)
        data_dict['mix_info_dict'] = mix_info

        return data_dict



class audio_video_spec_fullset_Dataset_Train(audio_video_spec_fullset_Dataset):
    def __init__(self, dataset_cfg):
        super().__init__(split='train', **dataset_cfg)

class audio_video_spec_fullset_Dataset_Valid(audio_video_spec_fullset_Dataset):
    def __init__(self, dataset_cfg):
        super().__init__(split='valid', **dataset_cfg)

class audio_video_spec_fullset_Dataset_Test(audio_video_spec_fullset_Dataset):
    def __init__(self, dataset_cfg):
        super().__init__(split='test', **dataset_cfg)



class audio_video_spec_fullset_Dataset_inpaint(audio_video_spec_fullset_Dataset):

    def __getitem__(self, idx):

        audio_name1 = self.data_list[idx]
        spec_npy_path1 = self.spec_list[idx]
        video_feat_path1 = self.feat_list[idx]
        video_path1 = self.video_list[idx]

        # Load the Spec and Feat:
        spec1, video_feat1 = self.load_spec_and_feat(spec_npy_path1, video_feat_path1)

        video_info_dict = {'audio_name1': audio_name1, 'audio_name2': "", 'video_path1': video_path1, 'video_path2': ""}
        mix_spec, mix_masked_spec, mix_video_feat, mix_info = self.mix_audio_and_feat(spec1=spec1, video_feat1=video_feat1, video_info_dict=video_info_dict)

        # print("mix spec shape:", mix_spec.shape)
        # print("mix video feat:", mix_video_feat.shape)
        data_dict = {}
        # data_dict['mix_spec'] = mix_spec[None].repeat(3, axis=0) # TODO：要把这里改掉，否则无法适应maa的autoencoder
        data_dict['mix_spec'] = mix_spec  # (80, 512)
        data_dict['hybrid_feat'] = {'mix_video_feat': mix_video_feat, 'mix_spec': mix_masked_spec}  # (32, 512)
        data_dict['mix_info_dict'] = mix_info

        return data_dict

    def mix_audio_and_feat(self, spec1=None, video_feat1=None, video_info_dict={}):
        """ Return Mix Spec and Mix video feat"""

        # spec1:
        if not self.fix_frames:
            start_idx = random.randint(0, self.sr * self.duration - self.truncate - 1)  # audio start
        else:
            start_idx = 0

        start_frame = int(self.fps * start_idx / self.sr)
        truncate_frame = int(self.fps * self.truncate / self.sr)

        # Spec Start & Truncate:
        spec_start = int(start_idx / self.hop_len)
        spec_truncate = int(self.truncate / self.hop_len)

        spec1 = spec1[:, spec_start: spec_start + spec_truncate]
        video_feat1 = video_feat1[start_frame: start_frame + truncate_frame]

        # Start masking frames:
        masked_spec = random.randint(1, int(spec_truncate * 0.5 // 16)) * 16  # 16帧的倍数，最多mask 50%
        masked_truncate = int(masked_spec * self.hop_len)
        masked_frame = int(self.fps * masked_truncate / self.sr)

        start_masked_idx = random.randint(0, self.truncate - masked_truncate - 1)
        start_masked_frame = int(self.fps * start_masked_idx / self.sr)
        start_masked_spec = int(start_masked_idx / self.hop_len)

        masked_spec1 = np.zeros((80, spec_truncate)).astype(np.float32)
        masked_spec1[:] = spec1[:]
        masked_spec1[:, start_masked_spec:start_masked_spec+masked_spec] = np.zeros((80, masked_spec))
        video_feat1[start_masked_frame:start_masked_frame+masked_frame, :] = np.zeros((masked_frame, 512))
        # info_dict:
        video_info_dict['video_time1'] = str(start_frame) + '_' + str(start_frame + truncate_frame)  # Start frame, end frame
        video_info_dict['video_time2'] = ""
        return spec1, masked_spec1, video_feat1, video_info_dict



class audio_video_spec_fullset_Dataset_inpaint_Train(audio_video_spec_fullset_Dataset_inpaint):
    def __init__(self, dataset_cfg):
        super().__init__(split='train', **dataset_cfg)

class audio_video_spec_fullset_Dataset_inpaint_Valid(audio_video_spec_fullset_Dataset_inpaint):
    def __init__(self, dataset_cfg):
        super().__init__(split='valid', **dataset_cfg)

class audio_video_spec_fullset_Dataset_inpaint_Test(audio_video_spec_fullset_Dataset_inpaint):
    def __init__(self, dataset_cfg):
        super().__init__(split='test', **dataset_cfg)



class audio_Dataset(torch.utils.data.Dataset):
    # Only Load audio dataset: for training Stage1: Audio Npy Dataset
    def __init__(self, split, dataset1, sr=22050, duration=10, truncate=220000, debug_num=False, fix_frames=False, hop_len=256):
        super().__init__()

        if split == "train":
            self.split = "Train"

        elif split == "valid" or split == 'test':
            self.split = "Test"

        # Default params:
        self.min_duration = 2
        self.sr = sr                # 22050
        self.duration = duration    # 10
        self.truncate = truncate    # 220000
        self.fix_frames = fix_frames
        self.hop_len = hop_len
        print("Fix Frames: {}".format(self.fix_frames))


        # Dataset1: (VGGSound)
        assert dataset1.dataset_name == "VGGSound"

        # spec_dir: spectrogram path

        # dataset1_spec_dir = os.path.join(dataset1.data_dir, "codec")
        dataset1_wav_dir = os.path.join(dataset1.wav_dir, "wav")

        split_txt_path = dataset1.split_txt_path
        with open(os.path.join(split_txt_path, '{}.txt'.format(self.split)), "r") as f:
            data_list1 = f.readlines()
            data_list1 = list(map(lambda x: x.strip(), data_list1))
            wav_list1 = list(map(lambda x: os.path.join(dataset1_wav_dir, x) + ".wav", data_list1))  # feat

        # Merge Data:
        self.data_list = data_list1
        self.wav_list = wav_list1

        assert len(self.data_list) == len(self.wav_list)

        shuffle_idx = np.random.permutation(np.arange(len(self.data_list)))
        self.data_list = [self.data_list[i] for i in shuffle_idx]
        self.wav_list = [self.wav_list[i] for i in shuffle_idx]

        if debug_num:
            self.data_list = self.data_list[:debug_num]
            self.wav_list = self.wav_list[:debug_num]

        print('Split: {}  Sample Num: {}'.format(split, len(self.data_list)))


    def __len__(self):
        return len(self.data_list)


    def load_spec_and_feat(self, wav_path):
        """Load audio spec and video feat"""
        try:
            wav_raw, sr = librosa.load(wav_path, sr=self.sr)                   # channel: 1
        except:
            print(f"corrupted wav: {wav_path}", flush=True)
            wav_raw = np.zeros((160000,), dtype=np.float32) # [T]

        wav_len = self.sr * self.duration
        if wav_raw.shape[0] < wav_len:
            wav_raw = np.tile(wav_raw, math.ceil(wav_len / wav_raw.shape[0]))
        wav_raw = wav_raw[:int(wav_len)]

        return wav_raw


    def mix_audio_and_feat(self, wav_raw1=None, video_info_dict={}, mode='single'):
        """ Return Mix Spec and Mix video feat"""
        if mode == "single":
            # spec1:
            if not self.fix_frames:
                start_idx = random.randint(0, self.sr * self.duration - self.truncate - 1)  # audio start
            else:
                start_idx = 0

            wav_start = start_idx
            wav_truncate = self.truncate
            wav_raw1 = wav_raw1[wav_start: wav_start + wav_truncate]

            return wav_raw1, video_info_dict

        elif mode == "concat":
            total_spec_len = int(self.truncate / self.hop_len)
            # Random Trucate len:
            spec1_truncate_len = random.randint(self.min_duration * self.sr // self.hop_len, total_spec_len - self.min_duration * self.sr // self.hop_len - 1)
            spec2_truncate_len = total_spec_len - spec1_truncate_len

            # Sample spec clip:
            spec_start1 = random.randint(0, total_spec_len - spec1_truncate_len - 1)
            spec_start2 = random.randint(0, total_spec_len - spec2_truncate_len - 1)
            spec_end1, spec_end2 = spec_start1 + spec1_truncate_len, spec_start2 + spec2_truncate_len

            # concat spec:
            return video_info_dict


    def __getitem__(self, idx):

        audio_name1 = self.data_list[idx]
        wav_path1 = self.wav_list[idx]
        # select other video:
        flag = False
        if random.uniform(0, 1) < -1:
            flag = True
            random_idx = idx
            while random_idx == idx:
                random_idx = random.randint(0, len(self.data_list)-1)
            audio_name2 = self.data_list[random_idx]
            spec_npy_path2 = self.spec_list[random_idx]
            wav_path2 = self.wav_list[random_idx]

        # Load the Spec and Feat:
        wav_raw1 = self.load_spec_and_feat(wav_path1)

        if flag:
            spec2, video_feat2 = self.load_spec_and_feat(spec_npy_path2, wav_path2)
            video_info_dict = {'audio_name1':audio_name1, 'audio_name2': audio_name2}
            mix_spec, mix_video_feat, mix_info = self.mix_audio_and_feat(video_info_dict, mode='concat')
        else:
            video_info_dict = {'audio_name1':audio_name1, 'audio_name2': ""}
            mix_wav, mix_info = self.mix_audio_and_feat(wav_raw1=wav_raw1, video_info_dict=video_info_dict, mode='single')

        data_dict = {}
        data_dict['mix_wav'] = mix_wav  # (131072,)
        data_dict['mix_info_dict'] = mix_info

        return data_dict


class audio_Dataset_Train(audio_Dataset):
    def __init__(self, dataset_cfg):
        super().__init__(split='train', **dataset_cfg)

class audio_Dataset_Test(audio_Dataset):
    def __init__(self, dataset_cfg):
        super().__init__(split='test', **dataset_cfg)

class audio_Dataset_Valid(audio_Dataset):
    def __init__(self, dataset_cfg):
        super().__init__(split='valid', **dataset_cfg)



class video_codec_Dataset(torch.utils.data.Dataset):
    # Only Load audio dataset: for training Stage1: Audio Npy Dataset
    def __init__(self, split, dataset1, sr=22050, duration=10, truncate=220000, fps=21.5, debug_num=False, fix_frames=False, hop_len=256):
        super().__init__()

        if split == "train":
            self.split = "Train"

        elif split == "valid" or split == 'test':
            self.split = "Test"

        # Default params:
        self.min_duration = 2
        self.fps = fps
        self.sr = sr                # 22050
        self.duration = duration    # 10
        self.truncate = truncate    # 220000
        self.fix_frames = fix_frames
        self.hop_len = hop_len
        print("Fix Frames: {}".format(self.fix_frames))


        # Dataset1: (VGGSound)
        assert dataset1.dataset_name == "VGGSound"

        # spec_dir: spectrogram path

        # dataset1_spec_dir = os.path.join(dataset1.data_dir, "codec")
        dataset1_feat_dir = os.path.join(dataset1.data_dir, "cavp")
        dataset1_wav_dir = os.path.join(dataset1.wav_dir, "wav")

        split_txt_path = dataset1.split_txt_path
        with open(os.path.join(split_txt_path, '{}.txt'.format(self.split)), "r") as f:
            data_list1 = f.readlines()
            data_list1 = list(map(lambda x: x.strip(), data_list1))
            wav_list1 = list(map(lambda x: os.path.join(dataset1_wav_dir, x) + ".wav", data_list1))  # feat
            feat_list1 = list(map(lambda x: os.path.join(dataset1_feat_dir, x) + ".npz", data_list1))  # feat

        # Merge Data:
        self.data_list = data_list1
        self.wav_list = wav_list1
        self.feat_list = feat_list1

        assert len(self.data_list) == len(self.wav_list)

        shuffle_idx = np.random.permutation(np.arange(len(self.data_list)))
        self.data_list = [self.data_list[i] for i in shuffle_idx]
        self.wav_list = [self.wav_list[i] for i in shuffle_idx]
        self.feat_list = [self.feat_list[i] for i in shuffle_idx]

        if debug_num:
            self.data_list = self.data_list[:debug_num]
            self.wav_list = self.wav_list[:debug_num]
            self.feat_list = self.feat_list[:debug_num]

        print('Split: {}  Sample Num: {}'.format(split, len(self.data_list)))


    def __len__(self):
        return len(self.data_list)


    def load_spec_and_feat(self, wav_path, video_feat_path):
        """Load audio spec and video feat"""
        try:
            wav_raw, sr = librosa.load(wav_path, sr=self.sr)                   # channel: 1
        except:
            print(f"corrupted wav: {wav_path}", flush=True)
            wav_raw = np.zeros((160000,), dtype=np.float32) # [T]

        try:
            video_feat = np.load(video_feat_path)['feat'].astype(np.float32)
        except:
            print(f"corrupted video: {video_feat_path}", flush=True)
            video_feat = np.load(os.path.join(os.path.dirname(video_feat_path), 'empty_vid.npz'))['feat'].astype(np.float32)

        wav_len = self.sr * self.duration
        if wav_raw.shape[0] < wav_len:
            wav_raw = np.tile(wav_raw, math.ceil(wav_len / wav_raw.shape[0]))
        wav_raw = wav_raw[:int(wav_len)]

        feat_len = self.fps * self.duration
        if video_feat.shape[0] < feat_len:
            video_feat = np.tile(video_feat, (math.ceil(feat_len / video_feat.shape[0]), 1))
        video_feat = video_feat[:int(feat_len)]

        return wav_raw, video_feat


    def mix_audio_and_feat(self, wav_raw1=None, video_feat1=None, video_info_dict={}, mode='single'):
        """ Return Mix Spec and Mix video feat"""
        if mode == "single":
            # spec1:
            if not self.fix_frames:
                start_idx = random.randint(0, self.sr * self.duration - self.truncate - 1)  # audio start
            else:
                start_idx = 0

            wav_start = start_idx
            wav_truncate = self.truncate
            wav_raw1 = wav_raw1[wav_start: wav_start + wav_truncate]

            start_frame = int(self.fps * start_idx / self.sr)
            truncate_frame = int(self.fps * self.truncate / self.sr)
            video_feat1 = video_feat1[start_frame: start_frame + truncate_frame]

            # info_dict:
            video_info_dict['video_time1'] = str(start_frame) + '_' + str(start_frame+truncate_frame)   # Start frame, end frame
            video_info_dict['video_time2'] = ""

            return wav_raw1, video_feat1, video_info_dict

        elif mode == "concat":
            total_spec_len = int(self.truncate / self.hop_len)
            # Random Trucate len:
            spec1_truncate_len = random.randint(self.min_duration * self.sr // self.hop_len, total_spec_len - self.min_duration * self.sr // self.hop_len - 1)
            spec2_truncate_len = total_spec_len - spec1_truncate_len

            # Sample spec clip:
            spec_start1 = random.randint(0, total_spec_len - spec1_truncate_len - 1)
            spec_start2 = random.randint(0, total_spec_len - spec2_truncate_len - 1)
            spec_end1, spec_end2 = spec_start1 + spec1_truncate_len, spec_start2 + spec2_truncate_len

            # concat spec:
            return video_info_dict


    def __getitem__(self, idx):

        audio_name1 = self.data_list[idx]
        wav_path1 = self.wav_list[idx]
        video_feat_path1 = self.feat_list[idx]
        # select other video:
        flag = False
        if random.uniform(0, 1) < -1:
            flag = True
            random_idx = idx
            while random_idx == idx:
                random_idx = random.randint(0, len(self.data_list)-1)
            audio_name2 = self.data_list[random_idx]
            wav_path2 = self.wav_list[random_idx]
            video_feat_path2 = self.feat_list[random_idx]

        # Load the Spec and Feat:
        wav_raw1, video_feat1 = self.load_spec_and_feat(wav_path1, video_feat_path1)

        if flag:
            wav_raw2, video_feat2 = self.load_spec_and_feat(wav_path2, video_feat_path2)
            video_info_dict = {'audio_name1':audio_name1, 'audio_name2': audio_name2}
            mix_spec, mix_video_feat, mix_info = self.mix_audio_and_feat(video_info_dict, mode='concat')
        else:
            video_info_dict = {'audio_name1':audio_name1, 'audio_name2': ""}
            mix_wav, mix_video_feat, mix_info = self.mix_audio_and_feat(wav_raw1=wav_raw1, video_feat1=video_feat1, video_info_dict=video_info_dict, mode='single')

        data_dict = {}
        data_dict['mix_wav'] = mix_wav  # (131072,)
        data_dict['mix_video_feat'] = mix_video_feat # (32, 512)
        data_dict['mix_info_dict'] = mix_info

        return data_dict


class video_codec_Dataset_Train(video_codec_Dataset):
    def __init__(self, dataset_cfg):
        super().__init__(split='train', **dataset_cfg)

class video_codec_Dataset_Test(video_codec_Dataset):
    def __init__(self, dataset_cfg):
        super().__init__(split='test', **dataset_cfg)

class video_codec_Dataset_Valid(video_codec_Dataset):
    def __init__(self, dataset_cfg):
        super().__init__(split='valid', **dataset_cfg)


class audio_video_spec_fullset_Audioset_Dataset(torch.utils.data.Dataset):
    # Only Load audio dataset: for training Stage1: Audio Npy Dataset
    def __init__(self, split, dataset1, dataset2, sr=22050, duration=10, truncate=220000,
                 fps=21.5, drop=0.0, fix_frames=False, hop_len=256):
        super().__init__()

        if split == "train":
            self.split = "Train"

        elif split == "valid" or split == 'test':
            self.split = "Test"

        # Default params:
        self.min_duration = 2
        self.sr = sr  # 22050
        self.duration = duration  # 10
        self.truncate = truncate  # 220000
        self.fps = fps
        self.fix_frames = fix_frames
        self.hop_len = hop_len
        self.drop = drop
        print("Fix Frames: {}".format(self.fix_frames))
        print("Use Drop: {}".format(self.drop))

        # Dataset1: (VGGSound)
        assert dataset1.dataset_name == "VGGSound"
        assert dataset2.dataset_name == "Audioset"

        # spec_dir: spectrogram path
        # feat_dir: CAVP feature path
        # video_dir: video path

        dataset1_spec_dir = os.path.join(dataset1.data_dir, "mel_maa2", "npy")
        dataset1_feat_dir = os.path.join(dataset1.data_dir, "cavp")
        split_txt_path = dataset1.split_txt_path
        with open(os.path.join(split_txt_path, '{}.txt'.format(self.split)), "r") as f:
            data_list1 = f.readlines()
            data_list1 = list(map(lambda x: x.strip(), data_list1))

            spec_list1 = list(map(lambda x: os.path.join(dataset1_spec_dir, x) + "_mel.npy", data_list1))  # spec
            feat_list1 = list(map(lambda x: os.path.join(dataset1_feat_dir, x) + ".npz", data_list1))  # feat

        if split == "train":
            dataset2_spec_dir = os.path.join(dataset2.data_dir, "mel")
            dataset2_feat_dir = os.path.join(dataset2.data_dir, "cavp_renamed")
            split_txt_path = dataset2.split_txt_path
            with open(os.path.join(split_txt_path, '{}.txt'.format(self.split)), "r") as f:
                data_list2 = f.readlines()
                data_list2 = list(map(lambda x: x.strip(), data_list2))

                spec_list2 = list(map(lambda x: os.path.join(dataset2_spec_dir, f'Y{x}') + "_mel.npy", data_list2))  # spec
                feat_list2 = list(map(lambda x: os.path.join(dataset2_feat_dir, x) + ".npz", data_list2))  # feat

            data_list1 += data_list2
            spec_list1 += spec_list2
            feat_list1 += feat_list2

        # Merge Data:
        self.data_list = data_list1 if self.split != "Test" else data_list1[:200]
        self.spec_list = spec_list1 if self.split != "Test" else spec_list1[:200]
        self.feat_list = feat_list1 if self.split != "Test" else feat_list1[:200]

        assert len(self.data_list) == len(self.spec_list) == len(self.feat_list)

        shuffle_idx = np.random.permutation(np.arange(len(self.data_list)))
        self.data_list = [self.data_list[i] for i in shuffle_idx]
        self.spec_list = [self.spec_list[i] for i in shuffle_idx]
        self.feat_list = [self.feat_list[i] for i in shuffle_idx]

        print('Split: {}  Sample Num: {}'.format(split, len(self.data_list)))

        # self.check(self.spec_list)

    def __len__(self):
        return len(self.data_list)

    def check(self, feat_list):
        from tqdm import tqdm
        for spec_path in tqdm(feat_list):
            mel = np.load(spec_path).astype(np.float32)
            if mel.shape[0] != 80:
                import ipdb
                ipdb.set_trace()



    def load_spec_and_feat(self, spec_path, video_feat_path):
        """Load audio spec and video feat"""
        spec_raw = np.load(spec_path).astype(np.float32)  # channel: 1
        if spec_raw.shape[0] != 80:
            print(f"corrupted mel: {spec_path}", flush=True)
            spec_raw = np.zeros((80, 625), dtype=np.float32)  # [C, T]

        p = np.random.uniform(0, 1)
        if p > self.drop:
            try:
                video_feat = np.load(video_feat_path)['feat'].astype(np.float32)
            except:
                print(f"corrupted video: {video_feat_path}", flush=True)
                video_feat = np.load(os.path.join(os.path.dirname(video_feat_path), 'empty_vid.npz'))['feat'].astype(np.float32)
        else:
            video_feat = np.load(os.path.join(os.path.dirname(video_feat_path), 'empty_vid.npz'))['feat'].astype(np.float32)

        spec_len = self.sr * self.duration / self.hop_len
        if spec_raw.shape[1] < spec_len:
            spec_raw = np.tile(spec_raw, math.ceil(spec_len / spec_raw.shape[1]))
        spec_raw = spec_raw[:, :int(spec_len)]

        feat_len = self.fps * self.duration
        if video_feat.shape[0] < feat_len:
            video_feat = np.tile(video_feat, (math.ceil(feat_len / video_feat.shape[0]), 1))
        video_feat = video_feat[:int(feat_len)]
        return spec_raw, video_feat

    def mix_audio_and_feat(self, spec1=None, spec2=None, video_feat1=None, video_feat2=None, video_info_dict={},
                           mode='single'):
        """ Return Mix Spec and Mix video feat"""
        if mode == "single":
            # spec1:
            if not self.fix_frames:
                start_idx = random.randint(0, self.sr * self.duration - self.truncate - 1)  # audio start
            else:
                start_idx = 0

            start_frame = int(self.fps * start_idx / self.sr)
            truncate_frame = int(self.fps * self.truncate / self.sr)

            # Spec Start & Truncate:
            spec_start = int(start_idx / self.hop_len)
            spec_truncate = int(self.truncate / self.hop_len)

            spec1 = spec1[:, spec_start: spec_start + spec_truncate]
            video_feat1 = video_feat1[start_frame: start_frame + truncate_frame]

            # info_dict:
            video_info_dict['video_time1'] = str(start_frame) + '_' + str(
                start_frame + truncate_frame)  # Start frame, end frame
            video_info_dict['video_time2'] = ""
            return spec1, video_feat1, video_info_dict

        elif mode == "concat":
            total_spec_len = int(self.truncate / self.hop_len)
            # Random Trucate len:
            spec1_truncate_len = random.randint(self.min_duration * self.sr // self.hop_len,
                                                total_spec_len - self.min_duration * self.sr // self.hop_len - 1)
            spec2_truncate_len = total_spec_len - spec1_truncate_len

            # Sample spec clip:
            spec_start1 = random.randint(0, total_spec_len - spec1_truncate_len - 1)
            spec_start2 = random.randint(0, total_spec_len - spec2_truncate_len - 1)
            spec_end1, spec_end2 = spec_start1 + spec1_truncate_len, spec_start2 + spec2_truncate_len

            # concat spec:
            spec1, spec2 = spec1[:, spec_start1: spec_end1], spec2[:, spec_start2: spec_end2]
            concat_audio_spec = np.concatenate([spec1, spec2], axis=1)

            # Concat Video Feat:
            start1_frame, truncate1_frame = int(self.fps * spec_start1 * self.hop_len / self.sr), int(
                self.fps * spec1_truncate_len * self.hop_len / self.sr)
            start2_frame, truncate2_frame = int(self.fps * spec_start2 * self.hop_len / self.sr), int(
                self.fps * self.truncate / self.sr) - truncate1_frame
            video_feat1, video_feat2 = video_feat1[start1_frame: start1_frame + truncate1_frame], video_feat2[
                                                                                                  start2_frame: start2_frame + truncate2_frame]
            concat_video_feat = np.concatenate([video_feat1, video_feat2])

            video_info_dict['video_time1'] = str(start1_frame) + '_' + str(
                start1_frame + truncate1_frame)  # Start frame, end frame
            video_info_dict['video_time2'] = str(start2_frame) + '_' + str(start2_frame + truncate2_frame)
            return concat_audio_spec, concat_video_feat, video_info_dict

    def __getitem__(self, idx):

        audio_name1 = self.data_list[idx]
        spec_npy_path1 = self.spec_list[idx]
        video_feat_path1 = self.feat_list[idx]

        # select other video:
        flag = False
        if random.uniform(0, 1) < -1:
            flag = True
            random_idx = idx
            while random_idx == idx:
                random_idx = random.randint(0, len(self.data_list) - 1)
            audio_name2 = self.data_list[random_idx]
            spec_npy_path2 = self.spec_list[random_idx]
            video_feat_path2 = self.feat_list[random_idx]

        # Load the Spec and Feat:
        spec1, video_feat1 = self.load_spec_and_feat(spec_npy_path1, video_feat_path1)

        if flag:
            spec2, video_feat2 = self.load_spec_and_feat(spec_npy_path2, video_feat_path2)
            video_info_dict = {'audio_name1': audio_name1, 'audio_name2': audio_name2}
            mix_spec, mix_video_feat, mix_info = self.mix_audio_and_feat(spec1, spec2, video_feat1, video_feat2, video_info_dict, mode='concat')
        else:
            video_info_dict = {'audio_name1': audio_name1, 'audio_name2': ""}
            mix_spec, mix_video_feat, mix_info = self.mix_audio_and_feat(spec1=spec1, video_feat1=video_feat1, video_info_dict=video_info_dict, mode='single')

        # print("mix spec shape:", mix_spec.shape)
        # print("mix video feat:", mix_video_feat.shape)
        data_dict = {}
        data_dict['mix_spec'] = mix_spec  # (80, 512)
        data_dict['mix_video_feat'] = mix_video_feat  # (32, 512)
        data_dict['mix_info_dict'] = mix_info

        return data_dict


class audio_video_spec_fullset_Audioset_Train(audio_video_spec_fullset_Audioset_Dataset):
    def __init__(self, dataset_cfg):
        super().__init__(split='train', **dataset_cfg)


class audio_video_spec_fullset_Audioset_Valid(audio_video_spec_fullset_Audioset_Dataset):
    def __init__(self, dataset_cfg):
        super().__init__(split='valid', **dataset_cfg)


class audio_video_spec_fullset_Audioset_Test(audio_video_spec_fullset_Audioset_Dataset):
    def __init__(self, dataset_cfg):
        super().__init__(split='test', **dataset_cfg)